package kb

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

// Reranker 语义重排客户端 (Cross-Encoder)
type Reranker struct {
	Endpoint string
	APIKey   string
	Model    string
	HTTP     *http.Client
}

// NewReranker 创建 Reranker 实例
func NewReranker(endpoint, key, model string) *Reranker {
	return &Reranker{
		Endpoint: normalizeAPIEndpoint(endpoint, "/rerank"),
		APIKey:   key,
		Model:    model,
		HTTP:     &http.Client{Timeout: 30 * time.Second},
	}
}

// Rerank 对候选文档列表依据查询进行相关性重排打分
// 返回与 docs 等长的一组打分，下标一一对应
func (r *Reranker) Rerank(ctx context.Context, query string, docs []string) ([]float64, error) {
	if r == nil || r.Endpoint == "" {
		return nil, fmt.Errorf("reranker not configured")
	}
	if len(docs) == 0 {
		return []float64{}, nil
	}

	body := map[string]any{
		"model":     r.Model,
		"query":     query,
		"documents": docs,
	}

	raw, err := postJSON(ctx, r.HTTP, r.Endpoint, r.APIKey, body)
	if err != nil {
		return nil, err
	}

	var resp struct {
		Results []struct {
			Index          int     `json:"index"`
			RelevanceScore float64 `json:"relevance_score"`
			Score          float64 `json:"score"`
		} `json:"results"`
	}

	if err := json.Unmarshal(raw, &resp); err != nil {
		return nil, fmt.Errorf("decode rerank: %s", formatHTTPBodyError(0, raw))
	}

	out := make([]float64, len(docs))
	for _, it := range resp.Results {
		s := it.RelevanceScore
		if s == 0 {
			s = it.Score
		}
		if it.Index >= 0 && it.Index < len(out) {
			out[it.Index] = s
		}
	}
	return out, nil
}