package vector

import (
	"context"
	"database/sql"
	"fmt"
	"log/slog"
	"strconv"
	"strings"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/embedding"

	_ "github.com/jackc/pgx/v5/stdlib" // Postgres driver, registers name "pgx"
)

// pgMaxOpenConns is the PostgreSQL connection pool size mandated by
// DevDocs/02-development-constraints.md §2.2 for the pgvector enhancement.
const pgMaxOpenConns = 25

// PGVectorIndex stores vectors in a PostgreSQL pgvector vector column and
// answers similarity with the cosine-distance operator <=> in SQL. It is the
// optional enhancement selected when DB_PG_URL is configured; when used,
// semantic search is genuinely backed by pgvector and does not fall back to
// keyword matching (DevDocs/design/03-记忆.md §3).
//
// The schema is prepared by NewPG (vector extension + vec_items table). Search
// issues ORDER BY vector <=> $1 LIMIT n, which returns exact cosine-distance-
// ranked results. With no ANN index the query planner does a sequential scan,
// which is the linear-scan behavior expected for milestone M4 (small datasets).
// An HNSW index (CREATE INDEX ... USING hnsw (vector vector_cosine_ops)) is the
// documented ANN upgrade for larger datasets and requires no change to Search.
type PGVectorIndex struct {
	db       *sql.DB
	dim      int
	embedder embedding.Embedder
	log      *slog.Logger
}

// NewPG connects to the PostgreSQL+pgvector database at pgURL using the pgx
// driver and prepares the schema: the vector extension and a vec_items table
// with a fixed-dimension vector column. It deliberately creates no ANN index:
// for milestone M4 the <=> query runs as an exact sequential scan, and an HNSW
// index is a documented, optional upgrade for larger data. A blank pgURL
// returns ErrInvalidPGURL, and a non-positive dim returns ErrInvalidDim, both
// before any network I/O. ctx controls the connection ping and schema setup.
//
// The returned index has no embedder wired in; use New to build a fully wired
// index, or wire one in before calling Embed.
func NewPG(ctx context.Context, pgURL string, dim int, log *slog.Logger) (*PGVectorIndex, error) {
	if strings.TrimSpace(pgURL) == "" {
		return nil, ErrInvalidPGURL
	}
	if dim <= 0 {
		return nil, ErrInvalidDim
	}
	log = sourcedLogger(log)

	db, err := sql.Open("pgx", pgURL)
	if err != nil {
		return nil, fmt.Errorf("memory/vector: open pg: %w", err)
	}
	db.SetMaxOpenConns(pgMaxOpenConns)
	if err := db.PingContext(ctx); err != nil {
		_ = db.Close()
		return nil, fmt.Errorf("memory/vector: ping pg: %w", err)
	}

	idx := &PGVectorIndex{db: db, dim: dim, log: log}
	if err := idx.setup(ctx); err != nil {
		_ = db.Close()
		return nil, err
	}
	log.Info("pgvector index ready", "dim", dim)
	return idx, nil
}

// setup creates the pgvector extension and the vec_items table idempotently.
// For milestone M4 no ANN index is created — <->/<=> queries fall back to an
// exact sequential scan, matching the linear-scan semantics of the SQLite path.
// To scale, add:
//
//	CREATE INDEX IF NOT EXISTS idx_vec_items_vector
//	    ON vec_items USING hnsw (vector vector_cosine_ops)
//
// which requires no change to Search. Each statement is parameterized where it
// takes input; the table DDL is static after dimension substitution.
func (p *PGVectorIndex) setup(ctx context.Context) error {
	statements := []string{
		"CREATE EXTENSION IF NOT EXISTS vector",
		fmt.Sprintf(`CREATE TABLE IF NOT EXISTS vec_items (
			id     TEXT PRIMARY KEY,
			vector vector(%d) NOT NULL,
			dim    INTEGER NOT NULL
		)`, p.dim),
	}
	for _, stmt := range statements {
		if _, err := p.db.ExecContext(ctx, stmt); err != nil {
			return fmt.Errorf("memory/vector: setup pg schema: %w", err)
		}
	}
	return nil
}

// withEmbedder wires the vector-producing embedder into the index. Used by New;
// tests may also call it directly.
func (p *PGVectorIndex) withEmbedder(e embedding.Embedder) { p.embedder = e }

// Dim returns the fixed vector dimension.
func (p *PGVectorIndex) Dim() int { return p.dim }

// Embed turns texts into vectors with the wired embedder, returning one Vector
// per input text aligned by index. With no embedder wired it returns
// ErrNoEmbedder.
func (p *PGVectorIndex) Embed(ctx context.Context, texts []string) ([]Vector, error) {
	if p.embedder == nil {
		return nil, ErrNoEmbedder
	}
	vecs, err := p.embedder.Embed(ctx, texts)
	if err != nil {
		return nil, fmt.Errorf("memory/vector: embed %d texts: %w", len(texts), err)
	}
	out := make([]Vector, len(vecs))
	for i, v := range vecs {
		out[i] = Vector(v)
	}
	return out, nil
}

// Add persists the (id, vector) pair, upserting on a repeated id. The vector is
// passed to Postgres as a vector literal bound to a $n::vector parameter. It
// rejects a blank id (ErrMissingID) and a vector whose length differs from the
// index dimension (ErrDimMismatch).
func (p *PGVectorIndex) Add(ctx context.Context, id string, v Vector) error {
	if strings.TrimSpace(id) == "" {
		return ErrMissingID
	}
	if len(v) != p.dim {
		return fmt.Errorf("memory/vector: %w (got %d, want %d)", ErrDimMismatch, len(v), p.dim)
	}
	_, err := p.db.ExecContext(ctx, `INSERT INTO vec_items (id, vector, dim)
		VALUES ($1, $2::vector, $3)
		ON CONFLICT (id) DO UPDATE SET vector = EXCLUDED.vector, dim = EXCLUDED.dim`,
		id, pgVectorLiteral(v), p.dim)
	if err != nil {
		return fmt.Errorf("memory/vector: add %q: %w", id, err)
	}
	return nil
}

// Search returns up to topK ids most similar to q, ranked by descending cosine
// similarity (1 - cosine distance), computed by pgvector in SQL. A non-positive
// topK returns an empty result; a query whose length differs from the index
// dimension returns ErrDimMismatch.
func (p *PGVectorIndex) Search(ctx context.Context, q Vector, topK int) ([]Hit, error) {
	if topK <= 0 {
		return []Hit{}, nil
	}
	if len(q) != p.dim {
		return nil, fmt.Errorf("memory/vector: %w (query %d, want %d)", ErrDimMismatch, len(q), p.dim)
	}
	rows, err := p.db.QueryContext(ctx, `SELECT id, 1.0 - (vector <=> $1::vector) AS score
		FROM vec_items
		ORDER BY vector <=> $1::vector
		LIMIT $2`, pgVectorLiteral(q), topK)
	if err != nil {
		return nil, fmt.Errorf("memory/vector: search pg: %w", err)
	}
	defer func() { _ = rows.Close() }()

	hits := make([]Hit, 0)
	for rows.Next() {
		var (
			id    string
			score float64
		)
		if err := rows.Scan(&id, &score); err != nil {
			return nil, fmt.Errorf("memory/vector: scan pg hit: %w", err)
		}
		hits = append(hits, Hit{ID: id, Score: score})
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("memory/vector: iterate pg hits: %w", err)
	}
	return hits, nil
}

// Close releases the underlying PostgreSQL connection pool.
func (p *PGVectorIndex) Close() error {
	if p.db == nil {
		return nil
	}
	err := p.db.Close()
	if err != nil {
		return fmt.Errorf("memory/vector: close pg: %w", err)
	}
	return nil
}

// pgVectorLiteral renders v in pgvector's text form "[a,b,c]" so it can be
// bound as a vector parameter ($n::vector). Values are formatted with enough
// precision to round-trip a float32.
func pgVectorLiteral(v Vector) string {
	var b strings.Builder
	b.WriteByte('[')
	for i, x := range v {
		if i > 0 {
			b.WriteByte(',')
		}
		b.WriteString(strconv.FormatFloat(float64(x), 'f', -1, 32))
	}
	b.WriteByte(']')
	return b.String()
}