package federated

// federated_test.go exercises the HARD spend gate (Phase 3) with in memory fakes for the ledger store and the
// audit seam, so the gate behavior is proven with NO database connected: an in budget release persists and is
// audited, an over budget release is refused and writes nothing, the accumulation is read from the persisted
// ledger never from the caller (never client trusted), and the Clip step is deterministic. The flag off
// refusal is covered in accountant_test.go.

import (
	"context"
	"encoding/json"
	"errors"
	"testing"

	"github.com/google/uuid"

	"github.com/missionworkforce/missionmatch/internal/platform/audit"
	"github.com/missionworkforce/missionmatch/internal/platform/config"
	"github.com/missionworkforce/missionmatch/internal/platform/database/dbgen"
)

// fakeLedger is an in memory LedgerStore. It records inserts and returns the latest as the running budget, so
// the gate composes across releases through the persisted (not caller supplied) accumulation.
type fakeLedger struct {
	rows       []dbgen.DpBudgetLedger
	insertErr  error // when set, Insert fails (to prove a failed write is not audited)
	latestErr  error // when set, LatestForKey fails
	insertSeen []dbgen.InsertDpBudgetLedgerParams
}

func (f *fakeLedger) LatestForKey(ctx context.Context, modelVersion string, tenantID *uuid.UUID) (dbgen.DpBudgetLedger, bool, error) {
	if f.latestErr != nil {
		return dbgen.DpBudgetLedger{}, false, f.latestErr
	}
	// Return the most recent row that matches the budget key, mirroring the SQL coalesce key match.
	var match *dbgen.DpBudgetLedger
	for i := range f.rows {
		r := f.rows[i]
		if r.ModelVersion != modelVersion {
			continue
		}
		if sameTenantKey(r.TenantID, tenantID) {
			match = &f.rows[i]
		}
	}
	if match == nil {
		return dbgen.DpBudgetLedger{}, false, nil
	}
	return *match, true, nil
}

func (f *fakeLedger) Insert(ctx context.Context, arg dbgen.InsertDpBudgetLedgerParams) (dbgen.DpBudgetLedger, error) {
	f.insertSeen = append(f.insertSeen, arg)
	if f.insertErr != nil {
		return dbgen.DpBudgetLedger{}, f.insertErr
	}
	row := dbgen.DpBudgetLedger{
		ID:                      uuid.New(),
		ModelVersion:            arg.ModelVersion,
		TenantID:                arg.TenantID,
		RdpOrders:               arg.RdpOrders,
		AccumulatedRdp:          arg.AccumulatedRdp,
		EpsilonAtDelta:          arg.EpsilonAtDelta,
		Delta:                   arg.Delta,
		ContributingTenantCount: arg.ContributingTenantCount,
	}
	f.rows = append(f.rows, row)
	return row, nil
}

func sameTenantKey(a, b *uuid.UUID) bool {
	if a == nil && b == nil {
		return true
	}
	if a == nil || b == nil {
		return false
	}
	return *a == *b
}

// fakeAuditor records the events the gate emits so the test asserts a release is audited as
// federated.dp.released with the right payload and no PII.
type fakeAuditor struct {
	events  []audit.Event
	failErr error
}

func (f *fakeAuditor) System(ctx context.Context, e audit.Event) error {
	if f.failErr != nil {
		return f.failErr
	}
	f.events = append(f.events, e)
	return nil
}

// newTestGate builds a gate over the fakes, flag ON, with a generous epsilon so the in budget path runs.
func newTestGate(store LedgerStore, auditW systemAuditor, epsilon, delta float64) *SpendGate {
	return &SpendGate{store: store, auditW: auditW, epsilon: epsilon, delta: delta, enabled: true}
}

// TestGateInBudgetReleasePersistsAndAudits proves an in budget release composes, persists exactly one ledger
// row through the store, and is audited as federated.dp.released with the model version, the contributing
// tenant count, and the converted epsilon and delta, and NO PII.
func TestGateInBudgetReleasePersistsAndAudits(t *testing.T) {
	store := &fakeLedger{}
	au := &fakeAuditor{}
	gate := newTestGate(store, au, 100, 1e-6)

	rel := Release{
		ModelVersion:            "v1.4.0",
		TenantID:                nil, // cross tenant cohort release
		ContributingTenantCount: 7,
		Gaussian:                GaussianRelease{NoiseSigma: 2, ClipNorm: 1},
	}
	eps, err := gate.Charge(context.Background(), rel)
	if err != nil {
		t.Fatalf("in budget Charge failed: %v", err)
	}
	if eps <= 0 {
		t.Fatalf("expected a positive epsilon, got %v", eps)
	}
	if len(store.rows) != 1 {
		t.Fatalf("expected exactly one persisted ledger row, got %d", len(store.rows))
	}
	if got := store.rows[0].ContributingTenantCount; got != 7 {
		t.Fatalf("contributing tenant count not recorded: got %d, want 7", got)
	}
	if len(au.events) != 1 {
		t.Fatalf("expected exactly one audit event, got %d", len(au.events))
	}
	ev := au.events[0]
	if ev.Action != "federated.dp.released" {
		t.Fatalf("audit action = %q, want federated.dp.released", ev.Action)
	}
	var payload map[string]any
	if err := json.Unmarshal(ev.NewValues, &payload); err != nil {
		t.Fatalf("audit payload not json: %v", err)
	}
	if payload["model_version"] != "v1.4.0" {
		t.Fatalf("audit payload model_version = %v, want v1.4.0", payload["model_version"])
	}
	// No PII, no candidate id, no raw vector in the payload.
	for _, banned := range []string{"candidate_id", "name", "email", "resume", "vector"} {
		if _, ok := payload[banned]; ok {
			t.Fatalf("audit payload leaked a banned field: %q", banned)
		}
	}
}

// TestGateComposesAcrossReleasesFromLedger proves the gate reads the PERSISTED accumulation (never a caller
// supplied one) and composes across releases: the epsilon after the second release is strictly greater than
// after the first, because the accumulation grew in the ledger.
func TestGateComposesAcrossReleasesFromLedger(t *testing.T) {
	store := &fakeLedger{}
	au := &fakeAuditor{}
	gate := newTestGate(store, au, 1e9, 1e-6) // huge budget so neither release is refused

	rel := Release{ModelVersion: "v1.4.0", ContributingTenantCount: 5, Gaussian: GaussianRelease{NoiseSigma: 2, ClipNorm: 1}}
	eps1, err := gate.Charge(context.Background(), rel)
	if err != nil {
		t.Fatalf("first Charge: %v", err)
	}
	eps2, err := gate.Charge(context.Background(), rel)
	if err != nil {
		t.Fatalf("second Charge: %v", err)
	}
	if eps2 <= eps1 {
		t.Fatalf("epsilon did not grow across composed releases read from the ledger: %v then %v", eps1, eps2)
	}
	if len(store.rows) != 2 {
		t.Fatalf("expected two persisted rows, got %d", len(store.rows))
	}
}

// TestGateOverBudgetRefusesAndWritesNothing proves an over budget release is refused with ErrBudgetExceeded,
// persists NO ledger row, and emits NO audit event. The budget is tiny so a single noisy release blows it.
func TestGateOverBudgetRefusesAndWritesNothing(t *testing.T) {
	store := &fakeLedger{}
	au := &fakeAuditor{}
	gate := newTestGate(store, au, 0.5, 1e-7) // tiny budget

	rel := Release{ModelVersion: "v1.4.0", ContributingTenantCount: 5, Gaussian: GaussianRelease{NoiseSigma: 0.1, ClipNorm: 1}} // huge RDP
	_, err := gate.Charge(context.Background(), rel)
	if !errors.Is(err, ErrBudgetExceeded) {
		t.Fatalf("expected ErrBudgetExceeded, got %v", err)
	}
	if len(store.rows) != 0 || len(store.insertSeen) != 0 {
		t.Fatalf("an over budget release wrote to the ledger: rows=%d inserts=%d", len(store.rows), len(store.insertSeen))
	}
	if len(au.events) != 0 {
		t.Fatalf("an over budget release was audited: %d events", len(au.events))
	}
}

// TestGateBudgetExhaustsAfterRepeatedReleases proves the budget is genuinely spent down: repeated in budget
// releases eventually trip the hard gate, after which every further release is refused and writes nothing.
func TestGateBudgetExhaustsAfterRepeatedReleases(t *testing.T) {
	store := &fakeLedger{}
	au := &fakeAuditor{}
	gate := newTestGate(store, au, 8, 1e-7) // moderate budget, reference Gboard band

	rel := Release{ModelVersion: "v1.4.0", ContributingTenantCount: 5, Gaussian: GaussianRelease{NoiseSigma: 1.5, ClipNorm: 1}}
	released, refused := 0, 0
	for i := 0; i < 1000; i++ {
		_, err := gate.Charge(context.Background(), rel)
		switch {
		case err == nil:
			released++
		case errors.Is(err, ErrBudgetExceeded):
			refused++
		default:
			t.Fatalf("unexpected error at iteration %d: %v", i, err)
		}
	}
	if released == 0 {
		t.Fatalf("budget never admitted a single release")
	}
	if refused == 0 {
		t.Fatalf("budget never exhausted, the hard gate never fired")
	}
	if len(store.rows) != released {
		t.Fatalf("persisted rows %d do not match released count %d", len(store.rows), released)
	}
	if len(au.events) != released {
		t.Fatalf("audit events %d do not match released count %d", len(au.events), released)
	}
}

// TestGateFailedWriteIsNotSilentlyAudited proves a ledger insert failure surfaces as an error and the release
// is NOT counted as audited beyond what the insert reached, so a write failure fails closed.
func TestGateFailedWriteIsNotSilentlyAudited(t *testing.T) {
	store := &fakeLedger{insertErr: errors.New("boom")}
	au := &fakeAuditor{}
	gate := newTestGate(store, au, 100, 1e-6)

	rel := Release{ModelVersion: "v1.4.0", ContributingTenantCount: 5, Gaussian: GaussianRelease{NoiseSigma: 2, ClipNorm: 1}}
	_, err := gate.Charge(context.Background(), rel)
	if err == nil {
		t.Fatalf("expected the ledger insert failure to surface")
	}
	if len(au.events) != 0 {
		t.Fatalf("a failed ledger write was still audited: %d events", len(au.events))
	}
}

// TestGateAuditFailureFailsClosed proves a release whose audit cannot be written returns an error (fails
// closed) rather than silently succeeding.
func TestGateAuditFailureFailsClosed(t *testing.T) {
	store := &fakeLedger{}
	au := &fakeAuditor{failErr: errors.New("audit down")}
	gate := newTestGate(store, au, 100, 1e-6)

	rel := Release{ModelVersion: "v1.4.0", ContributingTenantCount: 5, Gaussian: GaussianRelease{NoiseSigma: 2, ClipNorm: 1}}
	_, err := gate.Charge(context.Background(), rel)
	if err == nil {
		t.Fatalf("expected the audit failure to fail the release closed")
	}
}

// TestGateRemainingFlagOffIsZero proves Remaining returns zero headroom while the flag is off so nothing reads
// a live budget on a disabled deploy.
func TestGateRemainingFlagOffIsZero(t *testing.T) {
	cfg := config.Federated{DPEnabled: false, DPEpsilon: 10, DPDelta: 1e-7}
	gate := NewSpendGate(cfg, &fakeLedger{}, nil)
	rem, err := gate.Remaining(context.Background(), "v1.4.0", nil)
	if err != nil {
		t.Fatalf("Remaining flag off: %v", err)
	}
	if rem != 0 {
		t.Fatalf("flag off Remaining = %v, want 0", rem)
	}
}

// TestAggregatorEmitRoutesThroughGate proves the active aggregator's Emit goes through the gate (an in budget
// release persists and audits), confirming the Aggregator boundary is wired to the hard gate.
func TestAggregatorEmitRoutesThroughGate(t *testing.T) {
	store := &fakeLedger{}
	au := &fakeAuditor{}
	gate := newTestGate(store, au, 100, 1e-6)
	agg := &activeAggregator{gate: gate}

	_, err := agg.Emit(context.Background(), Release{ModelVersion: "v1.4.0", ContributingTenantCount: 5, Gaussian: GaussianRelease{NoiseSigma: 2, ClipNorm: 1}})
	if err != nil {
		t.Fatalf("Emit through the gate: %v", err)
	}
	if len(store.rows) != 1 || len(au.events) != 1 {
		t.Fatalf("Emit did not route through the gate: rows=%d events=%d", len(store.rows), len(au.events))
	}
}