package rag

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"sort"
	"time"
)

// rerankClient 调用 DashScope 文本重排（gte-rerank）。可选阶段：未配则跳过。
// DashScope 原生格式（非 OpenAI 兼容）：
//
//	POST {baseURL}  {model, input:{query, documents}, parameters:{top_n, return_documents:false}}
//	resp {output:{results:[{index, relevance_score}]}}
type rerankClient struct {
	baseURL string
	apiKey  string
	model   string
	hc      *http.Client
}

func newRerankClient(baseURL, apiKey, model string) *rerankClient {
	if baseURL == "" || model == "" {
		return nil
	}
	return &rerankClient{baseURL: baseURL, apiKey: apiKey, model: model, hc: &http.Client{Timeout: 20 * time.Second}}
}

func (r *rerankClient) ready() bool { return r != nil && r.baseURL != "" }

// rerank 用重排模型对候选重新打分排序，返回前 topN。出错时返回原序（降级）。
func (r *rerankClient) rerank(ctx context.Context, query string, hits []Hit, topN int) ([]Hit, error) {
	docs := make([]string, len(hits))
	for i, h := range hits {
		docs[i] = h.Text
	}
	body, _ := json.Marshal(map[string]any{
		"model": r.model,
		"input": map[string]any{"query": query, "documents": docs},
		"parameters": map[string]any{"top_n": topN, "return_documents": false},
	})
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, r.baseURL, bytes.NewReader(body))
	if err != nil {
		return hits, err
	}
	req.Header.Set("Content-Type", "application/json")
	if r.apiKey != "" {
		req.Header.Set("Authorization", "Bearer "+r.apiKey)
	}
	resp, err := r.hc.Do(req)
	if err != nil {
		return hits, err
	}
	defer resp.Body.Close()
	if resp.StatusCode >= 400 {
		buf := new(bytes.Buffer)
		_, _ = buf.ReadFrom(resp.Body)
		return hits, fmt.Errorf("rerank http %d: %s", resp.StatusCode, buf.String())
	}
	var out struct {
		Output struct {
			Results []struct {
				Index          int     `json:"index"`
				RelevanceScore float32 `json:"relevance_score"`
			} `json:"results"`
		} `json:"output"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
		return hits, err
	}
	res := make([]Hit, 0, len(out.Output.Results))
	for _, rr := range out.Output.Results {
		if rr.Index >= 0 && rr.Index < len(hits) {
			res = append(res, Hit{Text: hits[rr.Index].Text, Score: rr.RelevanceScore})
		}
	}
	if len(res) == 0 {
		return hits, nil
	}
	sort.Slice(res, func(i, j int) bool { return res[i].Score > res[j].Score })
	return res, nil
}