package ai

import (
	"context"
	"strings"
	"testing"

	"github.com/shopspring/decimal"

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

// TestWinProbabilityModelIntentRegisteredAndSafe proves the MAG win probability intent is wired safely into
// the closed Registry: it is registered, it is a model intent (Executor ExecutorModel), the whole Registry
// still validates (the MAG branch asserts its static shape), its OutputKind is the legal opportunity_score,
// its ModelKey is non empty, and its FeatureIntents are all registered names. This is the unit guarantee that
// a model intent cannot ship with an illegal shape.
func TestWinProbabilityModelIntentRegisteredAndSafe(t *testing.T) {
	intent, ok := IntentByName(IntentOpportunityWinProbability)
	if !ok {
		t.Fatalf("intent %q is not registered", IntentOpportunityWinProbability)
	}
	if intent.Executor != ExecutorModel {
		t.Fatalf("the win probability intent must be an ExecutorModel intent, got %q", intent.Executor)
	}
	if intent.Model == nil {
		t.Fatal("the win probability intent must carry a ModelSpec")
	}
	if intent.Model.OutputKind != OutputKindOpportunityScore {
		t.Errorf("the win probability intent must emit opportunity_score, got %q", intent.Model.OutputKind)
	}
	if intent.Model.ModelKey != WinProbabilityModelKey {
		t.Errorf("the win probability intent model key must be %q, got %q", WinProbabilityModelKey, intent.Model.ModelKey)
	}
	if strings.TrimSpace(intent.SQL) != "" {
		t.Error("a model intent must carry no reachable SQL")
	}
	// The whole Registry must still validate with the model intent and the new allowlist entries.
	if err := ValidateRegistry(); err != nil {
		t.Fatalf("ValidateRegistry failed with the win probability model intent: %v", err)
	}
	// The three MAG coefficient tables must be in the read allowlist (the reviewed decision).
	for _, tbl := range []string{"aedt_learned_models", "aedt_model_coefficients", "aedt_calibration_bins"} {
		if !ReadAllowlist[tbl] {
			t.Errorf("the MAG coefficient table %q must be in the AI ReadAllowlist", tbl)
		}
	}
}

// TestIsModelIntentClassification pins the routing classification: the win probability intent is a model
// intent, and it is NOT a plain SQL intent; the existing TAG and graph intents are NOT model intents. A
// misclassification would route a model intent through the raw SQL path (which rejects it) or run a SQL intent
// through the model executor.
func TestIsModelIntentClassification(t *testing.T) {
	if !IsModelIntent(IntentOpportunityWinProbability) {
		t.Fatal("the win probability intent must be a model intent")
	}
	// A model intent is not a SQL intent.
	intent, _ := IntentByName(IntentOpportunityWinProbability)
	if intent.IsSQLIntent() {
		t.Error("the win probability intent must not be a SQL intent")
	}
	// Existing intents are NOT model intents.
	for _, name := range []string{IntentBidWinRateBySegment, IntentLaborRateBand, "count_candidates_by_clearance"} {
		if IsModelIntent(name) {
			t.Errorf("%q is not a model intent, it must not classify as one", name)
		}
		it, _ := IntentByName(name)
		if !it.IsSQLIntent() {
			t.Errorf("%q must be a SQL intent", name)
		}
	}
	// An unknown name is not a model intent.
	if IsModelIntent("no_such_intent") {
		t.Error("an unknown name must not be a model intent")
	}
}

// TestValidateRegistryRejectsIllegalModelIntent proves the MAG branch of ValidateRegistry rejects a crafted
// model intent with an illegal OutputKind (a candidate score, the exact thing the enum forbids), an empty
// model key, an eeo feature intent, or a feature intent that is not in the Registry. It calls the branch
// directly on crafted intents so the closed Registry is never mutated. This is the load bearing proof that a
// candidate score is structurally impossible on the MAG surface.
func TestValidateRegistryRejectsIllegalModelIntent(t *testing.T) {
	registered := map[string]bool{IntentBidWinRateBySegment: true}
	bad := map[string]QueryIntent{
		"candidate score output kind": {
			Name:     "crafted_candidate_scorer",
			Executor: ExecutorModel,
			Model:    &ModelSpec{ModelKey: "k", OutputKind: "candidate_score", FeatureIntents: []string{IntentBidWinRateBySegment}},
		},
		"empty output kind": {
			Name:     "crafted_empty_kind",
			Executor: ExecutorModel,
			Model:    &ModelSpec{ModelKey: "k", OutputKind: "", FeatureIntents: []string{IntentBidWinRateBySegment}},
		},
		"empty model key": {
			Name:     "crafted_no_key",
			Executor: ExecutorModel,
			Model:    &ModelSpec{ModelKey: "", OutputKind: OutputKindOpportunityScore, FeatureIntents: []string{IntentBidWinRateBySegment}},
		},
		"eeo feature intent": {
			Name:     "crafted_eeo_feature",
			Executor: ExecutorModel,
			Model:    &ModelSpec{ModelKey: "k", OutputKind: OutputKindConfidence, FeatureIntents: []string{"eeo.candidate_demographics"}},
		},
		"unregistered feature intent": {
			Name:     "crafted_unregistered",
			Executor: ExecutorModel,
			Model:    &ModelSpec{ModelKey: "k", OutputKind: OutputKindDisplayOrder, FeatureIntents: []string{"no_such_intent"}},
		},
		"no feature intents": {
			Name:     "crafted_no_features",
			Executor: ExecutorModel,
			Model:    &ModelSpec{ModelKey: "k", OutputKind: OutputKindOpportunityScore, FeatureIntents: nil},
		},
		"reachable sql on a model intent": {
			Name:     "crafted_with_sql",
			Executor: ExecutorModel,
			SQL:      "SELECT 1 FROM candidates",
			Model:    &ModelSpec{ModelKey: "k", OutputKind: OutputKindOpportunityScore, FeatureIntents: []string{IntentBidWinRateBySegment}},
		},
		"nil model spec": {
			Name:     "crafted_nil_spec",
			Executor: ExecutorModel,
			Model:    nil,
		},
	}
	for name, it := range bad {
		if err := validateModelIntent(it, registered); err == nil {
			t.Errorf("validateModelIntent must reject %q", name)
		}
	}
	// A legal model intent passes.
	good := QueryIntent{
		Name:     "crafted_ok",
		Executor: ExecutorModel,
		Model:    &ModelSpec{ModelKey: "k", OutputKind: OutputKindOpportunityScore, FeatureIntents: []string{IntentBidWinRateBySegment}},
	}
	if err := validateModelIntent(good, registered); err != nil {
		t.Errorf("a legal model intent must pass: %v", err)
	}
	// The OutputKind enum forbids a candidate score at the validOutputKind check directly.
	for _, kind := range []string{"candidate_score", "match_score", "band", "total_score", ""} {
		if validOutputKind(kind) {
			t.Errorf("validOutputKind must reject %q; a candidate score is forbidden by construction", kind)
		}
	}
	for _, kind := range []string{OutputKindOpportunityScore, OutputKindDisplayOrder, OutputKindConfidence} {
		if !validOutputKind(kind) {
			t.Errorf("validOutputKind must accept the legal kind %q", kind)
		}
	}
}

// TestExecuteIntentRejectsModelIntentOnRawPath proves the defense in depth guard: the generic ExecuteIntent
// (the raw SQL path) hard rejects the win probability model intent, so a model intent can never run as raw
// SQL even if a future caller tried it. It needs no database (the rejection precedes any query).
func TestExecuteIntentRejectsModelIntentOnRawPath(t *testing.T) {
	intent, ok := IntentByName(IntentOpportunityWinProbability)
	if !ok {
		t.Fatalf("the win probability intent %q must be registered", IntentOpportunityWinProbability)
	}
	// A non nil db is required so the rejection is on the model kind, not the nil db guard. The db pointer's
	// Pool is never dereferenced because the model check precedes the query.
	_, err := ExecuteIntent(context.Background(), &tenant.DB{}, intent, []any{"agency", "naics"})
	if err == nil {
		t.Fatal("ExecuteIntent must reject the win probability model intent on the raw path, got nil error")
	}
	if !strings.Contains(err.Error(), "ExecuteModelIntent") {
		t.Errorf("the raw path rejection of a model intent must name ExecuteModelIntent, got %v", err)
	}
}

// TestScoreModelIntentGolden is the golden test of the MAG serving math: on a fixed champion (a known
// intercept and a known feature coefficient) and a fixed feature vector, the decimal log odds, the sigmoid,
// and the contribution trace are exactly as expected, and the contributions sum to the log odds. It proves
// the surfaced numbers are decimal and the additive explanation is faithful.
func TestScoreModelIntentGolden(t *testing.T) {
	// A champion with intercept logit(0.60) approx 0.4054651081 and a feature coefficient of 0.5. No
	// calibration bins, so the calibration is the identity map.
	intercept := decimal.RequireFromString("0.4054651081")
	model := championModel{
		ModelKey:   "test_model",
		Version:    1,
		OutputKind: OutputKindOpportunityScore,
		Intercept:  intercept,
		Coefficients: map[string]decimal.Decimal{
			"segment_base_rate": decimal.RequireFromString("0.5"),
		},
	}
	feats := []Contribution{
		{FeatureName: "segment_base_rate", Value: decimal.RequireFromString("0.4"), Rationale: "the segment base rate"},
	}

	value, trace := scoreModelIntent(model, feats)

	// The log odds is intercept + 0.5 * 0.4 = 0.4054651081 + 0.20 = 0.6054651081.
	// The trace: the intercept term (value 1, coefficient and contribution = intercept) plus the feature term
	// (value 0.4, coefficient 0.5, contribution 0.20).
	if len(trace) != 2 {
		t.Fatalf("the trace must carry the intercept plus one feature, got %d entries", len(trace))
	}
	if trace[0].FeatureName != interceptFeatureName {
		t.Errorf("the first trace entry must be the intercept, got %q", trace[0].FeatureName)
	}
	if !trace[0].Contribution.Equal(intercept) {
		t.Errorf("the intercept contribution must be %s, got %s", intercept, trace[0].Contribution)
	}
	wantFeatureContribution := decimal.RequireFromString("0.20")
	if !trace[1].Contribution.Equal(wantFeatureContribution) {
		t.Errorf("the feature contribution must be 0.20, got %s", trace[1].Contribution)
	}
	if !trace[1].Coefficient.Equal(decimal.RequireFromString("0.5")) {
		t.Errorf("the feature coefficient must be 0.5, got %s", trace[1].Coefficient)
	}

	// The contributions must sum to the log odds (the additive explanation is faithful).
	sum := decimal.Zero
	for _, c := range trace {
		sum = sum.Add(c.Contribution)
	}
	wantLogOdds := decimal.RequireFromString("0.6054651081")
	if !sum.Equal(wantLogOdds) {
		t.Errorf("the contributions must sum to the log odds %s, got %s", wantLogOdds, sum)
	}

	// The calibrated probability is sigmoid(0.6054651081) approx 0.6469. With no calibration bins the identity
	// map passes the raw through. The surfaced value is a decimal in [0, 1].
	if value.LessThan(decimal.Zero) || value.GreaterThan(decimal.NewFromInt(1)) {
		t.Fatalf("the probability must be in [0, 1], got %s", value)
	}
	// sigmoid(0.6054651081) = 1 / (1 + e^-0.6054651081). e^-0.6054651081 approx 0.54580, so p approx 0.64695.
	lo := decimal.RequireFromString("0.646")
	hi := decimal.RequireFromString("0.648")
	if value.LessThan(lo) || value.GreaterThan(hi) {
		t.Errorf("the calibrated probability must be near 0.647, got %s", value)
	}
}

// TestScoreModelIntentPriorOnlyReturnsBaseRate proves the prior only invariant: with every learned coefficient
// zero (the v0 seed), the model returns the sigmoid of the intercept alone (the calibrated base rate) no
// matter what the feature value is, because a zero coefficient makes the feature contribution zero. This is
// the correctness proof that the prior only v0 returns the shrunk base rate.
func TestScoreModelIntentPriorOnlyReturnsBaseRate(t *testing.T) {
	intercept := decimal.RequireFromString("0.4054651081") // logit(0.60)
	model := championModel{
		OutputKind:   OutputKindOpportunityScore,
		Intercept:    intercept,
		Coefficients: map[string]decimal.Decimal{"segment_base_rate": decimal.Zero}, // prior only: beta = 0
	}
	// Two very different feature values must yield the SAME probability, because the coefficient is zero.
	valueA, _ := scoreModelIntent(model, []Contribution{{FeatureName: "segment_base_rate", Value: decimal.RequireFromString("0.10")}})
	valueB, _ := scoreModelIntent(model, []Contribution{{FeatureName: "segment_base_rate", Value: decimal.RequireFromString("0.90")}})
	if !valueA.Equal(valueB) {
		t.Errorf("a prior only model (zero coefficient) must return the same base rate for any feature, got %s and %s", valueA, valueB)
	}
	// The base rate is sigmoid(logit(0.60)) = 0.60. Allow a small rounding tolerance.
	lo := decimal.RequireFromString("0.599")
	hi := decimal.RequireFromString("0.601")
	if valueA.LessThan(lo) || valueA.GreaterThan(hi) {
		t.Errorf("the prior only base rate must be near 0.60, got %s", valueA)
	}
}

// TestApplyCalibrationIdentityAndMap proves the calibration behavior: an empty bin set is the identity map (the
// raw passes through), and a bin that covers the raw remaps it to the bin's empirical rate. The result is
// always clamped to [0, 1].
func TestApplyCalibrationIdentityAndMap(t *testing.T) {
	raw := decimal.RequireFromString("0.647")
	// Empty bins: identity.
	if got := applyCalibration(nil, raw); !got.Equal(raw) {
		t.Errorf("an empty calibration set must be the identity map, got %s", got)
	}
	// A bin covering the raw remaps it.
	bins := []calibrationBin{{Lo: decimal.RequireFromString("0.60"), Hi: decimal.RequireFromString("0.70"), EmpiricalRate: decimal.RequireFromString("0.55")}}
	if got := applyCalibration(bins, raw); !got.Equal(decimal.RequireFromString("0.55")) {
		t.Errorf("a covering bin must remap the raw to its empirical rate, got %s", got)
	}
	// A raw outside every bin passes through (identity for the uncovered range).
	if got := applyCalibration(bins, decimal.RequireFromString("0.20")); !got.Equal(decimal.RequireFromString("0.20")) {
		t.Errorf("a raw outside every bin must pass through, got %s", got)
	}
}