package ai

import (
	"context"
	"errors"
	"fmt"
	"math"
	"sync"
	"sync/atomic"

	"github.com/jackc/pgx/v5/pgxpool"
	"github.com/pgvector/pgvector-go"

	"github.com/missionworkforce/missionmatch/internal/platform/embeddings"
	"github.com/missionworkforce/missionmatch/internal/platform/inference"
	"github.com/missionworkforce/missionmatch/internal/platform/tenant"
)

// ErrVectorModelMismatch means the stored vectors were not produced by the configured embed model (or
// the vector columns are the wrong width, or the store is not backfilled). Mixing embed models silently
// mis ranks every retrieval with no error, so the brain MUST refuse to serve until the store is re
// embedded with the configured model. The handler maps this to a typed 503 vector_model_mismatch.
var ErrVectorModelMismatch = errors.New("ai: stored vectors were not produced by the configured embed model")

// vectorMatchTolerance is the minimum cosine between a stored skill vector and a fresh embedding of the
// same skill name by the configured model. The same text under the same model, normalized, lands at
// cosine ~1.0; a DIFFERENT model (the all-minilm to bge corruption mode) lands far lower (the synonym
// and unrelated calibration showed a different meaning, never mind a different model, sits at 0.4 to
// 0.9), so 0.98 cleanly separates same model from a model swap while tolerating float drift.
const vectorMatchTolerance = 0.98

// VectorGate is the enforced re embed gate. It holds the boot time verdict on whether the stored vectors
// were produced by the configured embed model, and /ai/query consults it before serving. A genuine model
// mismatch closes the gate permanently (until a re embed and restart); a transient failure at boot (the
// plane was briefly unreachable) leaves the gate to self heal on the first query after the plane recovers,
// so a momentary outage does not dark the brain forever. The rest of the API never consults this gate.
type VectorGate struct {
	ok       atomic.Bool // fast path once verified, no lock
	mu       sync.Mutex
	mismatch bool   // a permanent negative verdict was reached
	reason   string // the last failure detail, for the log and the 503 message
	verify   func(ctx context.Context) error
}

// NewVectorGate builds a gate over a verification closure. A nil closure means no verifier is configured
// (the dev or test posture); such a gate is always open, so existing call sites that do not wire the
// check keep serving.
func NewVectorGate(verify func(ctx context.Context) error) *VectorGate {
	return &VectorGate{verify: verify}
}

// Check returns nil when the brain may serve, or a typed error when it may not: ErrVectorModelMismatch
// (a model mismatch or wrong width or unbackfilled store, a permanent closed gate) or an
// inference.ErrPlaneUnavailable wrap (the check could not reach the plane yet, retried on the next call).
func (g *VectorGate) Check(ctx context.Context) error {
	if g.ok.Load() {
		return nil
	}
	g.mu.Lock()
	defer g.mu.Unlock()
	if g.ok.Load() {
		return nil
	}
	if g.mismatch {
		return fmt.Errorf("%w: %s", ErrVectorModelMismatch, g.reason)
	}
	if g.verify == nil {
		g.ok.Store(true)
		return nil
	}
	err := g.verify(ctx)
	if err == nil {
		g.ok.Store(true)
		return nil
	}
	if errors.Is(err, ErrVectorModelMismatch) {
		g.mismatch = true
		g.reason = err.Error()
		return err
	}
	// Transient (the plane was unreachable for the re embed): stay unverified and surface a plane outage,
	// so the caller gets a typed 503 and the gate retries on the next query rather than darkening forever.
	g.reason = err.Error()
	return fmt.Errorf("%w: stored vector model not yet verified: %v", inference.ErrPlaneUnavailable, err)
}

// VerifyVectorModel is the boot self check. It asserts the vector column dimension for skills and
// candidates and confirms the stored skill vectors were produced by the configured embed model, by re
// embedding a sample of stored skill NAMES (the embedded text for a skill is its name, so it is
// re embeddable; a candidate's resume text is CUI and is never re embedded here, its column dimension is
// checked through the catalog instead). The skills table is global and shares the one backfill model with
// candidates, so a match here proves the configured model is the store's model.
func VerifyVectorModel(ctx context.Context, db *tenant.DB, client *inference.Client, dim int) error {
	if db == nil || db.Pool == nil {
		return errors.New("ai: vector model check needs a database")
	}
	if client == nil {
		return errors.New("ai: vector model check needs an inference client")
	}
	// Dimension contract through the catalog (the app role sees no candidate rows without a tenant, and
	// the column type carries the width), so a wrong width column is caught without reading any PII.
	if err := assertVectorColumnDim(ctx, db.Pool, "candidates", "resume_embedding", dim); err != nil {
		return err
	}
	if err := assertVectorColumnDim(ctx, db.Pool, "skills", "embedding", dim); err != nil {
		return err
	}
	// Model match on the global skills table. Re embed the SAME text the backfill embedded (the skill
	// name with its category), via the shared embeddings.SkillEmbedText, or a categorized skill would
	// look like a model mismatch when only the text differs.
	var texts []string
	var stored [][]float32
	rows, err := db.Pool.Query(ctx, "SELECT name, COALESCE(category,''), embedding FROM skills WHERE embedding IS NOT NULL ORDER BY id LIMIT 5")
	if err != nil {
		return fmt.Errorf("ai: sample stored skill vectors: %w", err)
	}
	defer rows.Close()
	for rows.Next() {
		var name, category string
		var v pgvector.Vector
		if err := rows.Scan(&name, &category, &v); err != nil {
			return err
		}
		texts = append(texts, embeddings.SkillEmbedText(name, category))
		stored = append(stored, v.Slice())
	}
	if err := rows.Err(); err != nil {
		return err
	}
	if len(texts) == 0 {
		return fmt.Errorf("%w: no backfilled skill embeddings to verify against (run the bge backfill before serving the brain)", ErrVectorModelMismatch)
	}
	res, err := client.Embed(ctx, texts)
	if err != nil {
		// Could not reach the plane to re embed: a transient condition, not a proven mismatch.
		return fmt.Errorf("ai: re embed for the vector model check: %w", err)
	}
	return vectorsMatchModel(texts, stored, res.Embeddings, dim)
}

// vectorsMatchModel is the pure comparison core (no DB, no plane), so the match and mismatch decision is
// unit testable. Every stored vector must be dim wide and cosine close to the fresh embedding of its name.
func vectorsMatchModel(names []string, stored, fresh [][]float32, dim int) error {
	if len(fresh) != len(stored) {
		return fmt.Errorf("ai: re embed returned %d vectors for %d stored", len(fresh), len(stored))
	}
	for i := range stored {
		if len(stored[i]) != dim {
			return fmt.Errorf("%w: stored skill %q vector width %d, expected %d", ErrVectorModelMismatch, names[i], len(stored[i]), dim)
		}
		if len(fresh[i]) != dim {
			return fmt.Errorf("ai: fresh embedding width %d, expected %d (plane or config error)", len(fresh[i]), dim)
		}
		c := cosineSim(stored[i], fresh[i])
		if c < vectorMatchTolerance {
			return fmt.Errorf("%w: stored vector for skill %q matches the configured model at cosine %.4f, below tolerance %.2f (the store was embedded by a different model, re embed before serving)", ErrVectorModelMismatch, names[i], c, vectorMatchTolerance)
		}
	}
	return nil
}

// assertVectorColumnDim reads the pgvector column dimension from the catalog (pgvector stores the
// dimension directly in atttypmod, with no offset) and asserts it equals dim.
func assertVectorColumnDim(ctx context.Context, pool *pgxpool.Pool, table, column string, dim int) error {
	var typmod int
	err := pool.QueryRow(ctx,
		`SELECT a.atttypmod FROM pg_attribute a
		   JOIN pg_class c ON c.oid = a.attrelid
		   JOIN pg_namespace n ON n.oid = c.relnamespace
		  WHERE n.nspname = 'public' AND c.relname = $1 AND a.attname = $2 AND NOT a.attisdropped`,
		table, column).Scan(&typmod)
	if err != nil {
		return fmt.Errorf("ai: read %s.%s vector dimension: %w", table, column, err)
	}
	if typmod != dim {
		return fmt.Errorf("%w: column %s.%s is vector(%d), expected vector(%d)", ErrVectorModelMismatch, table, column, typmod, dim)
	}
	return nil
}

// cosineSim computes the true cosine (it divides by both norms), so it is correct whether or not the
// stored vector was normalized at backfill.
func cosineSim(a, b []float32) float64 {
	var dot, na, nb float64
	for i := range a {
		dot += float64(a[i]) * float64(b[i])
		na += float64(a[i]) * float64(a[i])
		nb += float64(b[i]) * float64(b[i])
	}
	if na == 0 || nb == 0 {
		return 0
	}
	return dot / (math.Sqrt(na) * math.Sqrt(nb))
}