package federated

// federated.go defines the cross tenant federated aggregator boundary (Clip, Noise, Emit) and the HARD
// server side spend gate that backs Charge with the dp_budget_ledger. Everything here is DEAD CODE unless
// MM_FEDERATED_DP_ENABLED is set: the constructor for an active aggregator refuses to build while the flag is
// false, so there is no live emission path today, and the noise sampler is an activation blocker that returns
// an error rather than sampling. The package has no non test caller; cmd/teach-aggregate in Phase 5 will be
// the gated caller.
//
// THE GATE IS NEVER CLIENT TRUSTED. Charge reads the PERSISTED accumulated RDP for a (model_version,
// tenant_id) budget key from the ledger, composes the candidate release into it, reconverts to epsilon at the
// configured delta, and REFUSES the release (writing nothing) if it would exceed MM_FEDERATED_DP_EPSILON.
// Only an in budget release persists a new ledger row, and it does so through the MIGRATOR pool, because
// dp_budget_ledger is a global migrator owned table the app role can only SELECT. No client supplied epsilon,
// no client supplied accumulation, ever enters the decision.
//
// AUDIT RESPONSIBILITY (documented, load bearing). A successful, in budget release is audited as
// federated.dp.released through audit.Writer.System (NULL org, a cross tenant global event) BY THE GATE
// ITSELF, inside Charge, after the ledger write succeeds. The gate owns the audit so a release can never be
// persisted without its audit record. The event carries the model version, the contributing tenant count,
// the converted epsilon and delta, and the worst case accumulated RDP, never any PII, candidate id, or raw
// per tenant vector. If the audit cannot be written the release fails closed.

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"math"
	"strconv"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
	"github.com/shopspring/decimal"

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

// ErrDisabled is returned by the active aggregator constructor when MM_FEDERATED_DP_ENABLED is false. It is
// the guard that keeps the whole emission path dead in tenant LIVE: an active aggregator cannot even be
// constructed while the flag is off.
var ErrDisabled = errors.New("federated: MM_FEDERATED_DP_ENABLED is false, the active aggregator is disabled")

// Aggregator is the cross tenant federated aggregator boundary. It is the seam for the central DP v1
// implementation and its TEE or secure aggregation upgrade. The three steps are the DP-FedEmb pattern: Clip
// each per tenant aggregate to a fixed L2 sensitivity, Noise the clipped aggregate with calibrated Gaussian
// noise, then Emit the noised aggregate through the spend gate (which charges the budget and persists the
// release). Today no implementation can be constructed while the flag is off, so none of these run.
type Aggregator interface {
	// Clip bounds a per tenant aggregate vector to L2 norm clipNorm (the sensitivity C), returning the
	// clipped vector. Clipping is the deterministic step that establishes the per tenant sensitivity the
	// accountant calibrates noise against. It does not require the secure sampler and is safe to compute.
	Clip(vec []float64, clipNorm float64) []float64
	// Noise adds calibrated Gaussian noise to a clipped vector. THIS IS THE ACTIVATION BLOCKER: it must use a
	// Mironov safe discrete or snapped Gaussian sampler behind the FIPS boundary, which is not implemented, so
	// the only implementation returns ErrNotImplemented. The naive float Gaussian is intentionally absent so
	// it can never run on an activated path.
	Noise(vec []float64, noiseSigma float64) ([]float64, error)
	// Emit charges the release against the budget through the gate and, only if in budget, persists the
	// release to the ledger and audits it. It returns the epsilon at the configured delta after the release,
	// or ErrBudgetExceeded if the budget is spent.
	Emit(ctx context.Context, release Release) (epsilonNow float64, err error)
}

// Release is one cross tenant differential privacy release to emit. ModelVersion is the shared model version
// the release contributes to (pinned from engine.ModelVersion by the caller, never a literal). TenantID is
// nil for a cross tenant cohort release, which has no single owning tenant. ContributingTenantCount is the
// number of live non test distinct organizations in the cohort, sourced from the production flag, and is
// recorded so the (epsilon, delta) claim is always reported with its cohort size. Gaussian carries the noise
// multiplier inputs the accountant charges. The noised aggregate vector itself is NOT carried here because
// the gate accounts and audits counts only; the vector lives in the shared model artifact, never the ledger
// or the audit log.
type Release struct {
	ModelVersion            string
	TenantID                *uuid.UUID
	ContributingTenantCount int
	Gaussian                GaussianRelease
}

// LedgerStore is the migrator pool seam the gate reads the running budget through and writes a release
// through. The live implementation is backed by the MIGRATOR pool (the app role has SELECT only on
// dp_budget_ledger); a test substitutes a fake in memory store. The package depends on this interface, not
// the concrete dbgen.DBTX, so the gate is testable without a database and the migrator pool is supplied by
// the caller (the cmd/teach-aggregate two pool wiring in Phase 5), never the app role.
type LedgerStore interface {
	// LatestForKey returns the most recent ledger row for the (modelVersion, tenantID) budget key, carrying
	// the running accumulated RDP composed so far. It returns ok false when the budget is fresh (no prior
	// release), in which case the accumulation starts at zero.
	LatestForKey(ctx context.Context, modelVersion string, tenantID *uuid.UUID) (row dbgen.DpBudgetLedger, ok bool, err error)
	// Insert persists one in budget release row and returns it. The gate calls this ONLY after it has
	// confirmed the release stays within budget, so a persisted row is by construction an in budget release.
	Insert(ctx context.Context, arg dbgen.InsertDpBudgetLedgerParams) (dbgen.DpBudgetLedger, error)
}

// systemAuditor is the off pool audit seam the gate writes federated.dp.released through. *audit.Writer
// satisfies it via System, so production wiring passes a real writer unchanged; a test substitutes a fake to
// assert the release is audited without a database. The gate depends on this interface, not the concrete
// writer, for the same testability reason the ledger uses LedgerStore.
type systemAuditor interface {
	System(ctx context.Context, e audit.Event) error
}

// SpendGate is the hard server side differential privacy spend gate. It is the ledger backed Accountant: each
// Charge reads the persisted accumulated RDP for the release's budget key, composes the candidate release,
// reconverts to epsilon at the configured delta, refuses (writing nothing) if over the budget, and otherwise
// persists the new release row through the migrator pool and audits it as federated.dp.released. The gate is
// the only path that writes the ledger, and it is never client trusted: the accumulation is read from the
// persisted ledger, never supplied by a caller.
type SpendGate struct {
	store   LedgerStore   // the migrator pool seam (dp_budget_ledger reads and writes)
	auditW  systemAuditor // audit.Writer.System for the federated.dp.released event (NULL org, global)
	epsilon float64       // the epsilon budget target, from config (MM_FEDERATED_DP_EPSILON), never a constant
	delta   float64       // the delta target, from config (MM_FEDERATED_DP_DELTA), never a constant
	enabled bool          // MM_FEDERATED_DP_ENABLED; false makes every Charge a no op refusal (dead code guard)
}

// NewSpendGate builds the hard spend gate from configuration. The flag, epsilon, and delta come from the
// Federated config block, never from code constants, so a privacy officer signs off and tunes them without a
// rebuild. store is the MIGRATOR pool seam (never the app pool, which cannot write the global ledger) and
// auditW writes the federated.dp.released system event (pass a real *audit.Writer in production). The gate is
// built regardless of the flag, but when the flag is false every Charge refuses with ErrDisabled and writes
// nothing, so the emission path stays dead. A caller that wants a hard construction guard should use
// NewActiveAggregator, which refuses to build at all while the flag is off.
func NewSpendGate(cfg config.Federated, store LedgerStore, auditW *audit.Writer) *SpendGate {
	return &SpendGate{
		store:   store,
		auditW:  auditW,
		epsilon: cfg.DPEpsilon,
		delta:   cfg.DPDelta,
		enabled: cfg.DPEnabled,
	}
}

// chargeRelease is the shared budget decision for one release. It loads the persisted accumulation, composes
// the candidate release, and returns the post release accountant plus the per order accumulation, or
// ErrBudgetExceeded. It performs NO write and NO audit, so it is the pure decision the Emit path and any
// dry run share. The flag guard lives in Charge and Emit, not here.
func (g *SpendGate) chargeRelease(ctx context.Context, release Release) (acct *rdpAccountant, epsilonNow float64, err error) {
	acct, err = g.loadAccountant(ctx, release.ModelVersion, release.TenantID)
	if err != nil {
		return nil, 0, err
	}
	epsilonNow, err = acct.Charge(release.Gaussian)
	if err != nil {
		return nil, 0, err
	}
	return acct, epsilonNow, nil
}

// loadAccountant builds an RDP accountant seeded from the persisted accumulation for the (modelVersion,
// tenantID) budget key. A fresh budget (no prior release) starts at zero; a resumed budget decodes the per
// order accumulation out of the stored rdp_orders jsonb so composition continues across process restarts.
func (g *SpendGate) loadAccountant(ctx context.Context, modelVersion string, tenantID *uuid.UUID) (*rdpAccountant, error) {
	prior, ok, err := g.store.LatestForKey(ctx, modelVersion, tenantID)
	if err != nil {
		return nil, fmt.Errorf("federated: read budget ledger: %w", err)
	}
	acct, err := newRDPAccountant(g.epsilon, g.delta, nil)
	if err != nil {
		return nil, err
	}
	if ok {
		seed, err := decodeRDPOrders(prior.RdpOrders, acct.orders)
		if err != nil {
			return nil, err
		}
		acct, err = newRDPAccountant(g.epsilon, g.delta, seed)
		if err != nil {
			return nil, err
		}
	}
	return acct, nil
}

// Charge is the hard spend gate. While the flag is off it refuses with ErrDisabled and writes nothing (the
// dead code guard). While the flag is on it reads the persisted accumulated RDP for the budget key, composes
// the candidate release, refuses with ErrBudgetExceeded (writing nothing) if over budget, and otherwise
// persists the new release row through the migrator pool and audits federated.dp.released through
// audit.Writer.System. The persisted accumulated_rdp and epsilon_at_delta are shopspring decimals so the
// audit record is bit stable; the grid math ran in float64. A persisted row is by construction an in budget,
// already audited release.
func (g *SpendGate) Charge(ctx context.Context, release Release) (epsilonNow float64, err error) {
	if !g.enabled {
		return 0, ErrDisabled
	}
	acct, epsilonNow, err := g.chargeRelease(ctx, release)
	if err != nil {
		return 0, err
	}

	ordersJSON, err := encodeRDPOrders(acct.accumulatedRDPDecimal())
	if err != nil {
		return 0, err
	}
	row, err := g.store.Insert(ctx, dbgen.InsertDpBudgetLedgerParams{
		ModelVersion:            release.ModelVersion,
		TenantID:                release.TenantID,
		RdpOrders:               ordersJSON,
		AccumulatedRdp:          acct.worstRDPDecimal(),
		EpsilonAtDelta:          acct.epsilonDecimal(),
		Delta:                   decimal.NewFromFloat(g.delta),
		ContributingTenantCount: int32(release.ContributingTenantCount),
	})
	if err != nil {
		return 0, fmt.Errorf("federated: persist budget release: %w", err)
	}

	if err := g.auditRelease(ctx, row); err != nil {
		// Fail closed: a release the audit cannot record is treated as failed even though the row persisted,
		// because the dp_budget_ledger is append only and the release is already counted against the budget.
		return 0, fmt.Errorf("federated: audit dp release: %w", err)
	}
	return epsilonNow, nil
}

// Remaining returns the epsilon headroom left on a budget key. It is a read only budget check (no write, no
// audit) so it is safe to call regardless of the flag, but it returns zero headroom while the flag is off so
// nothing reads a live budget on a disabled deploy. tenantID is nil for the cross tenant cohort key.
func (g *SpendGate) Remaining(ctx context.Context, modelVersion string, tenantID *uuid.UUID) (float64, error) {
	if !g.enabled {
		return 0, nil
	}
	acct, err := g.loadAccountant(ctx, modelVersion, tenantID)
	if err != nil {
		return 0, err
	}
	return acct.Remaining(), nil
}

// auditRelease writes the federated.dp.released system event (NULL org, a cross tenant global event). It
// carries the model version, the contributing tenant count, the converted epsilon and delta, and the worst
// case accumulated RDP, never any PII, candidate id, or raw vector.
func (g *SpendGate) auditRelease(ctx context.Context, row dbgen.DpBudgetLedger) error {
	payload := map[string]any{
		"model_version":             row.ModelVersion,
		"contributing_tenant_count": row.ContributingTenantCount,
		"epsilon_at_delta":          row.EpsilonAtDelta.String(),
		"delta":                     row.Delta.String(),
		"accumulated_rdp":           row.AccumulatedRdp.String(),
	}
	body, err := json.Marshal(payload)
	if err != nil {
		return err
	}
	return g.auditW.System(ctx, audit.Event{
		Action:     "federated.dp.released",
		EntityType: "dp_budget_ledger",
		EntityID:   row.ID.String(),
		NewValues:  body,
	})
}

// activeAggregator is the only Aggregator implementation. It is constructed ONLY when the flag is on (see
// NewActiveAggregator), so its emission path is unreachable in tenant LIVE. Clip is deterministic and safe;
// Noise is the activation blocker; Emit routes through the spend gate.
type activeAggregator struct {
	gate *SpendGate
}

// NewActiveAggregator is the HARD construction guard for the emission path. It refuses to build an active
// aggregator (returning ErrDisabled) while MM_FEDERATED_DP_ENABLED is false, so no live emission path can
// even be assembled in tenant LIVE. When the flag is on it returns an aggregator whose Emit routes through
// the spend gate; its Noise still returns ErrNotImplemented until the Mironov safe sampler is written, so an
// activated path cannot actually sample noise either, only assemble the pieces.
func NewActiveAggregator(cfg config.Federated, gate *SpendGate) (Aggregator, error) {
	if !cfg.DPEnabled {
		return nil, ErrDisabled
	}
	return &activeAggregator{gate: gate}, nil
}

// Clip bounds vec to L2 norm clipNorm. It is the deterministic sensitivity step and does not sample noise, so
// it is safe to compute even on a flag off path (no constructor returns this type while the flag is off, but
// the method itself touches no randomness). A non positive clipNorm or a vector already inside the ball is
// returned unchanged.
func (a *activeAggregator) Clip(vec []float64, clipNorm float64) []float64 {
	if clipNorm <= 0 {
		return vec
	}
	var sumSq float64
	for _, v := range vec {
		sumSq += v * v
	}
	norm := math.Sqrt(sumSq)
	if norm <= clipNorm || norm == 0 {
		return vec
	}
	scale := clipNorm / norm
	out := make([]float64, len(vec))
	for i, v := range vec {
		out[i] = v * scale
	}
	return out
}

// Noise is the HARD activation blocker. It returns ErrNotImplemented and samples nothing. The Mironov safe
// discrete or snapped Gaussian sampler (the only acceptable activated sampler, behind the FIPS boundary) is a
// documented TODO; the naive float Gaussian is intentionally absent from the whole package so it can never
// run. See docs/AI_RESEARCH.md gotcha 7 (floating point DP attacks, Mironov 2012).
//
// TODO activation blocker: implement a Mironov safe discrete or snapped Gaussian sampler that draws only from
// crypto/rand under GOFIPS140 (no new crypto dependency, in module primitives only) before the flag is ever
// flipped on. Until then this returns an error and no noise is added.
func (a *activeAggregator) Noise(vec []float64, noiseSigma float64) ([]float64, error) {
	return nil, ErrNotImplemented
}

// Emit routes the release through the hard spend gate, which charges the budget, persists an in budget
// release, and audits it. It returns the epsilon at the configured delta after the release or
// ErrBudgetExceeded.
func (a *activeAggregator) Emit(ctx context.Context, release Release) (float64, error) {
	return a.gate.Charge(ctx, release)
}

// migratorLedgerStore backs LedgerStore with a migrator scoped dbgen.DBTX (the migrator pool). It is a thin
// adapter so the gate depends on the interface, not the concrete DBTX, and a test substitutes a fake. The
// caller passes the MIGRATOR pool here, never the app pool, because the app role cannot write the global
// dp_budget_ledger.
type migratorLedgerStore struct{ db dbgen.DBTX }

// NewMigratorLedgerStore wraps a migrator scoped dbgen.DBTX as the ledger store. The caller passes the
// migrator pool, never the app pool.
func NewMigratorLedgerStore(migrator dbgen.DBTX) LedgerStore {
	return migratorLedgerStore{db: migrator}
}

func (s migratorLedgerStore) LatestForKey(ctx context.Context, modelVersion string, tenantID *uuid.UUID) (dbgen.DpBudgetLedger, bool, error) {
	row, err := dbgen.New(s.db).GetLatestDpBudgetLedger(ctx, dbgen.GetLatestDpBudgetLedgerParams{
		ModelVersion: modelVersion,
		TenantID:     tenantID,
	})
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return dbgen.DpBudgetLedger{}, false, nil
		}
		return dbgen.DpBudgetLedger{}, false, err
	}
	return row, true, nil
}

func (s migratorLedgerStore) Insert(ctx context.Context, arg dbgen.InsertDpBudgetLedgerParams) (dbgen.DpBudgetLedger, error) {
	return dbgen.New(s.db).InsertDpBudgetLedger(ctx, arg)
}

// encodeRDPOrders serializes the per order accumulated RDP decimal map to the rdp_orders jsonb, keyed by the
// Renyi order as a string so the json object is stable and the audit record reproduces exactly.
func encodeRDPOrders(m map[float64]decimal.Decimal) ([]byte, error) {
	obj := make(map[string]string, len(m))
	for order, v := range m {
		obj[strconv.FormatFloat(order, 'g', -1, 64)] = v.String()
	}
	return json.Marshal(obj)
}

// decodeRDPOrders parses the rdp_orders jsonb back into a per order accumulated RDP slice index aligned with
// the grid, so composition resumes from a persisted ledger row. An order in the grid that is absent from the
// stored map seeds to zero (a fresh order); the float64 grid math then resumes from the persisted values.
func decodeRDPOrders(raw []byte, orders []float64) ([]float64, error) {
	out := make([]float64, len(orders))
	if len(raw) == 0 {
		return out, nil
	}
	var obj map[string]string
	if err := json.Unmarshal(raw, &obj); err != nil {
		return nil, fmt.Errorf("federated: decode rdp_orders: %w", err)
	}
	for i, order := range orders {
		key := strconv.FormatFloat(order, 'g', -1, 64)
		s, ok := obj[key]
		if !ok {
			continue
		}
		v, err := strconv.ParseFloat(s, 64)
		if err != nil {
			return nil, fmt.Errorf("federated: parse rdp_orders value: %w", err)
		}
		out[i] = v
	}
	return out, nil
}