package kb

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"

	"ai-operator/internal/embedding"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

type PostgresRepository struct{ pool *pgxpool.Pool }

func NewPostgresRepository(pool *pgxpool.Pool) *PostgresRepository {
	return &PostgresRepository{pool: pool}
}

func (r *PostgresRepository) Health(ctx context.Context) (Health, error) {
	h := Health{}
	if err := r.pool.Ping(ctx); err != nil {
		return h, err
	}
	h.DBReachable = true
	_ = r.pool.QueryRow(ctx, `SELECT EXISTS (SELECT 1 FROM pg_extension WHERE extname='vector')`).Scan(&h.VectorExtension)
	var mig int
	_ = r.pool.QueryRow(ctx, `SELECT count(*) FROM information_schema.tables WHERE table_name IN ('schema_migrations','knowledge_documents','knowledge_chunks','knowledge_ingest_runs','knowledge_search_logs')`).Scan(&mig)
	h.MigrationsApplied = mig == 5
	_ = r.pool.QueryRow(ctx, `SELECT count(*) FROM knowledge_documents`).Scan(&h.Documents)
	_ = r.pool.QueryRow(ctx, `SELECT count(*) FROM knowledge_chunks`).Scan(&h.Chunks)
	return h, nil
}

func (r *PostgresRepository) UpsertDocument(ctx context.Context, doc Document, chunks []Chunk) error {
	tx, err := r.pool.Begin(ctx)
	if err != nil {
		return err
	}
	defer tx.Rollback(ctx)
	docID := stableUUID(doc.ExternalID)
	meta, _ := json.Marshal(doc.Metadata)
	_, err = tx.Exec(ctx, `INSERT INTO knowledge_documents(id,external_id,title,language,region_code,category,source_type,source_uri,source_hash,status,content,question,short_answer,full_answer,keywords,source,metadata,priority,updated_at) VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,now()) ON CONFLICT(external_id) DO UPDATE SET title=EXCLUDED.title, language=EXCLUDED.language, region_code=EXCLUDED.region_code, category=EXCLUDED.category, source_type=EXCLUDED.source_type, source_uri=EXCLUDED.source_uri, source_hash=EXCLUDED.source_hash, status=EXCLUDED.status, content=EXCLUDED.content, question=EXCLUDED.question, short_answer=EXCLUDED.short_answer, full_answer=EXCLUDED.full_answer, keywords=EXCLUDED.keywords, source=EXCLUDED.source, metadata=EXCLUDED.metadata, priority=EXCLUDED.priority, updated_at=now()`, docID, doc.ExternalID, doc.Title, doc.Language, doc.RegionCode, doc.Category, doc.SourceType, doc.SourceURI, doc.SourceHash, doc.Status, doc.Content, nullable(doc.Question), nullable(doc.ShortAnswer), nullable(doc.FullAnswer), nonNilKeywords(doc.Keywords), nullable(doc.Source), meta, doc.Priority)
	if err != nil {
		return err
	}
	if _, err = tx.Exec(ctx, `DELETE FROM knowledge_chunks WHERE document_id=$1`, docID); err != nil {
		return err
	}
	for _, ch := range chunks {
		chID := stableUUID(fmt.Sprintf("%s:%d:%s", doc.ExternalID, ch.ChunkIndex, ch.ContentHash))
		cm, _ := json.Marshal(ch.Metadata)
		_, err = tx.Exec(ctx, `INSERT INTO knowledge_chunks(id,document_id,chunk_index,language,region_code,title,content,content_hash,token_estimate,char_count,embedding,embedding_model,embedding_provider,metadata,tsv) VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11::vector,$12,$13,$14,to_tsvector('simple',$7))`, chID, docID, ch.ChunkIndex, ch.Language, ch.RegionCode, ch.Title, ch.Content, ch.ContentHash, ch.TokenEstimate, ch.CharCount, embedding.VectorLiteral(ch.Embedding), ch.EmbeddingModel, ch.EmbeddingProvider, cm)
		if err != nil {
			return err
		}
	}
	return tx.Commit(ctx)
}

func (r *PostgresRepository) Search(ctx context.Context, req SearchRequest, queryEmbedding []float32) ([]SearchResult, error) {
	limit := req.Limit
	if limit <= 0 {
		limit = 5
	}
	if limit > 10 {
		limit = 10
	}
	langs := []string{req.Language}
	regions := []string{req.RegionCode}
	if req.IncludeGlobal {
		regions = append(regions, "global")
	}
	vec := embedding.VectorLiteral(queryEmbedding)
	rows, err := r.pool.Query(ctx, `WITH q AS (SELECT plainto_tsquery('simple', $1) query) SELECT d.id::text, c.id::text, c.title, c.content, c.language, c.region_code, d.category, coalesce(d.source_uri,''), coalesce(d.source,''), (1 - ((c.embedding <=> $2::vector)/2.0))::float8 AS vector_score, ts_rank_cd(c.tsv, q.query)::float8 AS keyword_score, CASE WHEN c.region_code=$3 THEN 1.0 ELSE 0.6 END AS region_boost FROM knowledge_chunks c JOIN knowledge_documents d ON d.id=c.document_id, q WHERE d.status='published' AND c.language=ANY($4) AND c.region_code=ANY($5) AND (d.valid_from IS NULL OR d.valid_from<=now()) AND (d.valid_to IS NULL OR d.valid_to>=now()) AND (c.tsv @@ q.query OR c.content ILIKE '%' || $1 || '%' OR c.embedding IS NOT NULL) ORDER BY ((0.65*(1-((c.embedding <=> $2::vector)/2.0))) + (0.25*ts_rank_cd(c.tsv,q.query)) + (0.10*CASE WHEN c.region_code=$3 THEN 1.0 ELSE 0.6 END)) DESC LIMIT $6`, sanitizeQuery(req.Query), vec, req.RegionCode, langs, regions, limit)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []SearchResult
	for rows.Next() {
		var r0 SearchResult
		var source string
		var v, k, boost float64
		if err := rows.Scan(&r0.DocumentID, &r0.ChunkID, &r0.Title, &r0.Content, &r0.Language, &r0.RegionCode, &r0.Category, &r0.SourceURI, &source, &v, &k, &boost); err != nil {
			return nil, err
		}
		r0.SourceLanguage = r0.Language
		r0.VectorScore = v
		r0.KeywordScore = k
		r0.Score = 0.65*v + 0.25*k + 0.10*boost
		r0.Citation = Citation{DocumentTitle: r0.Title, SourceURI: r0.SourceURI, ChunkID: r0.ChunkID, RegionCode: r0.RegionCode, Language: r0.Language, Source: source}
		if r0.Score >= req.MinScore {
			out = append(out, r0)
		}
	}
	_ = r.logSearch(ctx, req, out, false)
	return out, rows.Err()
}

func (r *PostgresRepository) logSearch(ctx context.Context, req SearchRequest, results []SearchResult, fallback bool) error {
	top := 0.0
	if len(results) > 0 {
		top = results[0].Score
	}
	_, err := r.pool.Exec(ctx, `INSERT INTO knowledge_search_logs(id,call_id,query,language,region_code,result_count,top_score,cross_language_fallback_used) VALUES($1,$2,$3,$4,$5,$6,$7,$8)`, stableUUID(fmt.Sprintf("%s:%s:%d", req.CallID, req.Query, len(results))), nullable(req.CallID), truncate(req.Query, 1000), req.Language, req.RegionCode, len(results), top, fallback)
	return err
}

func nullable(s string) any {
	if strings.TrimSpace(s) == "" {
		return nil
	}
	return s
}
func sanitizeQuery(s string) string {
	return strings.Map(func(r rune) rune {
		if r < 32 {
			return -1
		}
		return r
	}, truncate(strings.TrimSpace(s), 1000))
}
func truncate(s string, n int) string {
	r := []rune(s)
	if len(r) > n {
		return string(r[:n])
	}
	return s
}
func stableUUID(s string) string {
	h := ContentHash(s)
	return fmt.Sprintf("%s-%s-%s-%s-%s", h[0:8], h[8:12], h[12:16], h[16:20], h[20:32])
}

var _ pgx.Tx

func nonNilKeywords(v []string) []string {
	if v == nil {
		return []string{}
	}
	return v
}