package services

import (
	"fmt"
	"math"
	"sort"

	"github.com/sujit-baniya/fiber-boilerplate/pkg/models"
)

// resultadoRAG empareja un chunk con su similitud a la consulta, para poder ordenar.
type resultadoRAG struct {
	chunk     models.UmindChunk
	similitud float64
}

// BuscarConocimiento embebe la consulta del usuario y devuelve los topK
// fragmentos más parecidos de la base de conocimiento del agente, por
// similitud coseno calculada en memoria. Sin pgvector por ahora: para el
// volumen de un piloto (un agente, unos cientos de chunks) esto es
// suficientemente rápido; si el volumen crece, se reemplaza por una consulta
// pgvector sin cambiar la firma de esta función.
func BuscarConocimiento(agenteID uint, consulta string, topK int) ([]models.UmindChunk, error) {
	if topK <= 0 {
		topK = 4
	}

	ai, err := models.GetUmindEmbeddingsConfig()
	if err != nil {
		return nil, err
	}
	consultaVec, err := GenerarEmbedding(ai, consulta)
	if err != nil {
		return nil, fmt.Errorf("no se pudo generar el embedding de la consulta: %w", err)
	}

	chunks, err := models.GetUmindChunksByAgente(agenteID)
	if err != nil {
		return nil, err
	}
	if len(chunks) == 0 {
		return nil, nil
	}

	resultados := make([]resultadoRAG, 0, len(chunks))
	for _, c := range chunks {
		vec, err := models.EmbeddingFromJSON(c.EmbeddingJSON)
		if err != nil || len(vec) == 0 {
			continue
		}
		resultados = append(resultados, resultadoRAG{chunk: c, similitud: similitudCoseno(consultaVec, vec)})
	}

	sort.Slice(resultados, func(i, j int) bool { return resultados[i].similitud > resultados[j].similitud })

	if topK > len(resultados) {
		topK = len(resultados)
	}
	out := make([]models.UmindChunk, topK)
	for i := 0; i < topK; i++ {
		out[i] = resultados[i].chunk
	}
	return out, nil
}

// similitudCoseno calcula la similitud coseno entre dos vectores del mismo
// tamaño. Vectores de tamaño distinto (embeddings de modelos diferentes) dan 0.
func similitudCoseno(a, b []float32) float64 {
	if len(a) != len(b) || len(a) == 0 {
		return 0
	}
	var punto, normaA, normaB float64
	for i := range a {
		ai, bi := float64(a[i]), float64(b[i])
		punto += ai * bi
		normaA += ai * ai
		normaB += bi * bi
	}
	if normaA == 0 || normaB == 0 {
		return 0
	}
	return punto / (math.Sqrt(normaA) * math.Sqrt(normaB))
}