package eino

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

	"github.com/cloudwego/eino/schema"
	"github.com/gogf/gf/v2/frame/g"
)

// DashScopeReranker 通义百炼 Rerank 精排（Cross-Encoder）
type DashScopeReranker struct {
	httpClient *http.Client
}

func NewDashScopeReranker() *DashScopeReranker {
	return &DashScopeReranker{
		httpClient: &http.Client{
			Timeout: 10 * time.Second,
		},
	}
}

// Rerank 对文档进行精排（Cross-Encoder 核心）
func (d *DashScopeReranker) Rerank(ctx context.Context, query string, docs []*schema.Document) ([]*schema.Document, error) {
	if len(docs) == 0 {
		return docs, nil
	}

	// 官方必过 URL
	url := "https://dashscope.aliyuncs.com/api/v1/services/rerank/text-rerank/text-rerank"
	apiKey := g.Cfg().MustGet(ctx, "eino.rerank.apiKey").String()
	model := g.Cfg().MustGet(ctx, "eino.rerank.model").String()

	documents := make([]string, len(docs))
	for i, doc := range docs {
		documents[i] = doc.Content
	}

	reqBody := map[string]any{
		"model": model,
		"input": map[string]any{
			"query":     query,
			"documents": documents,
		},
		"parameters": map[string]any{
			"top_n": len(docs),
		},
	}
	bs, _ := json.Marshal(reqBody)

	req, _ := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(bs))
	req.Header.Set("Authorization", "Bearer "+apiKey)
	req.Header.Set("Content-Type", "application/json")

	resp, err := d.httpClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		body, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("rerank api error: status=%d, body=%s", resp.StatusCode, string(body))
	}

	// 解析结果
	var result struct {
		Output struct {
			Results []struct {
				Index          int     `json:"index"`
				RelevanceScore float64 `json:"relevance_score"`
			} `json:"results"`
		} `json:"output"`
	}

	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return nil, err
	}

	// 按分数排序
	type scoredDoc struct {
		doc   *schema.Document
		score float64
	}
	scored := make([]scoredDoc, len(docs))
	for i, doc := range docs {
		scored[i] = scoredDoc{doc: doc, score: 0}
	}
	for _, res := range result.Output.Results {
		scored[res.Index].score = res.RelevanceScore
	}

	// 分数从高到低排序
	for i := 0; i < len(scored); i++ {
		for j := i + 1; j < len(scored); j++ {
			if scored[j].score > scored[i].score {
				scored[i], scored[j] = scored[j], scored[i]
			}
		}
	}

	// 输出最终排好的文档
	ranked := make([]*schema.Document, 0, len(scored))
	for _, s := range scored {
		s.doc.MetaData["rerank_score"] = s.score
		ranked = append(ranked, s.doc)
	}

	return ranked, nil
}