package eino

import (
	"context"
	"errors"

	"github.com/cloudwego/eino/callbacks"
	"github.com/cloudwego/eino/components/embedding"
	"github.com/cloudwego/eino/components/retriever"
	"github.com/cloudwego/eino/schema"
	"github.com/gogf/gf/v2/util/gconv"
	"github.com/pgvector/pgvector-go"
)

type PGVectorRetrieverConfig struct {
	Embedder     embedding.Embedder
	DefaultTopK  int
	DefaultIndex string
}

type PGVectorRetriever struct {
	embedder embedding.Embedder
	topK     int
	index    string
}

func NewPGVectorRetriever(config *PGVectorRetrieverConfig) (*PGVectorRetriever, error) {
	if config.Embedder == nil {
		return nil, errors.New("embedder is required")
	}
	if config.DefaultTopK <= 0 {
		config.DefaultTopK = 5
	}

	return &PGVectorRetriever{
		embedder: config.Embedder,
		topK:     config.DefaultTopK,
		index:    config.DefaultIndex,
	}, nil
}

func (r *PGVectorRetriever) Retrieve(ctx context.Context, query string, opts ...retriever.Option) ([]*schema.Document, error) {

	// 1. 处理公共 Option（官方标准写法）
	options := &retriever.Options{
		Index:     &r.index,
		TopK:      &r.topK,
		Embedding: r.embedder,
	}
	options = retriever.GetCommonOptions(options, opts...)

	// 2. 回调（官方标准）
	ctx = callbacks.OnStart(ctx, &retriever.CallbackInput{
		Query: query,
		TopK:  *options.TopK,
	})

	// 3. 执行检索
	docs, err := r.doRetrieve(ctx, query, options)
	if err != nil {
		callbacks.OnError(ctx, err)
		return nil, err
	}

	// 4. 完成回调
	callbacks.OnEnd(ctx, &retriever.CallbackOutput{
		Docs: docs,
	})

	return docs, nil
}

func (r *PGVectorRetriever) doRetrieve(ctx context.Context, query string, opts *retriever.Options) ([]*schema.Document, error) {

	// 1. 生成向量
	vectors, err := opts.Embedding.EmbedStrings(ctx, []string{query})
	if err != nil {
		return nil, err
	}
	if len(vectors) == 0 {
		return nil, errors.New("empty query vector")
	}

	queryVec := pgvector.NewVector(vectors[0])
	topK := *opts.TopK

	// 2. PG 向量相似度检索 SQL
	sql := `
	SELECT id, content, dataset_id, document_id,
	       vector <-> ? AS distance
	FROM document_chunk
	ORDER BY distance ASC
	LIMIT ?
`

	// 3. 查询
	rows, err := dao.DocumentChunk.GetDB().GetAll(ctx, sql, queryVec, topK)
	if err != nil {
		return nil, err
	}

	// 4. 转为 Eino Document
	docs := make([]*schema.Document, 0, len(rows))
	for _, row := range rows {
		docs = append(docs, &schema.Document{
			ID:      gconv.String(row["id"]),
			Content: gconv.String(row["content"]),
			Metadata: map[string]any{
				"dataset_id":  row["dataset_id"],
				"document_id": row["document_id"],
				"distance":    row["distance"],
			},
		})
	}

	return docs, nil
}