package magstore

import (
	"context"
	"fmt"

	"github.com/jackc/pgx/v5/pgxpool"

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

// TimeToFillModelKey is the model_key of the discrete time survival forecast of how long a staffing requirement
// stays open before it is filled (migration 0039 seeds the prior only v1 constant baseline hazard).
const TimeToFillModelKey = "requirement_time_to_fill"

// HorizonMonths is the fixed hazard curve length (one coefficient per month). The survival is computed over
// these months; a fill within horizon question beyond it uses the full curve.
const HorizonMonths = 12

// hazardFeatureName is the stored coefficient name for a month's hazard (hazard_month_1 .. hazard_month_12), the
// same names migration 0039 seeds and the serve path reads.
func hazardFeatureName(month int) string {
	return fmt.Sprintf("hazard_month_%d", month)
}

// ReadRequirementDurations reads every requirement's time to fill observation POOLED across tenants: the whole
// months from creation to the FIRST approved slate member (the fill event), or, when no approval exists, the
// months from creation to now, right censored (still open). Months are ceil of the elapsed time, clamped to at
// least 1 and at most maxPeriod. A soft deleted requirement is excluded. Requirements are joined to their
// slates' approved members by the composite org key, so a cross tenant join cannot happen under the migrator.
func ReadRequirementDurations(ctx context.Context, pool *pgxpool.Pool, maxPeriod int) ([]learn.Duration, error) {
	if maxPeriod < 1 {
		maxPeriod = HorizonMonths
	}
	const q = `
SELECT
    GREATEST(1, LEAST($1::int, CEIL(EXTRACT(EPOCH FROM (COALESCE(f.filled_at, now()) - r.created_at)) / (30.0 * 86400))))::int AS periods,
    (f.filled_at IS NOT NULL) AS filled
  FROM requirements r
  LEFT JOIN LATERAL (
      SELECT min(sm.decided_at) AS filled_at
        FROM slate_members sm
        JOIN slates s ON s.organization_id = sm.organization_id AND s.id = sm.slate_id
       WHERE s.requirement_id = r.id AND sm.decision = 'approved'
  ) f ON true
 WHERE r.deleted_at IS NULL
 ORDER BY r.id`
	rows, err := pool.Query(ctx, q, maxPeriod)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []learn.Duration
	for rows.Next() {
		var d learn.Duration
		if err := rows.Scan(&d.Periods, &d.Filled); err != nil {
			return nil, err
		}
		out = append(out, d)
	}
	return out, rows.Err()
}

// TimeToFillTrainResult summarizes a time to fill training run.
type TimeToFillTrainResult struct {
	Skipped      bool
	SkipReason   string
	Version      int
	Requirements int
	Filled       int
	Hazard       []float64 // the fitted per month hazard curve
}

// TrainTimeToFill reads the pooled requirement durations, estimates the empirical discrete time hazard curve,
// and writes a new CHALLENGER version of the time to fill model. It refuses to fit below the requirement floor
// or when no requirement has ever been filled (no fill signal to learn from), leaving the prior only baseline
// champion in force. It NEVER flips champion. The hazard is stored as one coefficient per month; a discrete
// hazard model carries no calibration bins.
func TrainTimeToFill(ctx context.Context, pool *pgxpool.Pool, modelKey string, minRequirements int, dryRun bool) (TimeToFillTrainResult, error) {
	durations, err := ReadRequirementDurations(ctx, pool, HorizonMonths)
	if err != nil {
		return TimeToFillTrainResult{}, err
	}
	res := TimeToFillTrainResult{Requirements: len(durations)}
	for _, d := range durations {
		if d.Filled {
			res.Filled++
		}
	}
	if minRequirements <= 0 {
		minRequirements = 30
	}
	if res.Requirements < minRequirements {
		res.Skipped = true
		res.SkipReason = fmt.Sprintf("%d requirements is below the floor of %d", res.Requirements, minRequirements)
		return res, nil
	}
	if res.Filled == 0 {
		res.Skipped = true
		res.SkipReason = "no requirement has been filled yet, so there is no fill hazard to learn"
		return res, nil
	}

	res.Hazard = learn.EmpiricalHazard(durations, HorizonMonths)

	if dryRun {
		return res, nil
	}

	names := make([]string, HorizonMonths)
	for m := 1; m <= HorizonMonths; m++ {
		names[m-1] = hazardFeatureName(m)
	}
	notes := fmt.Sprintf("Challenger empirical discrete time hazard fit by cmd/train-model over %d requirements (%d filled) pooled across tenants. OutputKind opportunity_score expressing P(filled within a horizon), never a candidate score. champion=false until cmd/eval-model passes the calibration gate and a human promotes it.", res.Requirements, res.Filled)
	version, werr := WriteChallenger(ctx, pool, modelKey, Challenger{
		Family:       "discrete_hazard",
		OutputKind:   "opportunity_score",
		FeatureNames: names,
		Intercept:    0,
		Coeffs:       res.Hazard,
		Notes:        notes,
	})
	if werr != nil {
		return TimeToFillTrainResult{}, werr
	}
	res.Version = version
	return res, nil
}

// HazardCurve extracts a loaded model's per month hazard vector in month order (hazard_month_1 .. _12), so a
// stored time to fill model (champion or challenger) can form its survival curve and answer a fill within
// horizon question.
func HazardCurve(m Model) []float64 {
	hz := make([]float64, HorizonMonths)
	for month := 1; month <= HorizonMonths; month++ {
		hz[month-1] = m.Coeffs[hazardFeatureName(month)]
	}
	return hz
}

// FillWithinHorizon returns the model's probability that a requirement is filled within `horizon` months: one
// minus the cumulative survival at the horizon. It is the served opportunity level number; the per month hazard
// curve is the interpretable evidence.
func FillWithinHorizon(m Model, horizon int) float64 {
	return learn.FillWithin(HazardCurve(m), horizon)
}

// TimeToFillEval is the held out CALIBRATION comparison of the champion and a challenger time to fill model at a
// horizon. A baseline hazard is one curve for every requirement, so the meaningful quality metric is
// calibration: does the predicted fill within horizon probability match the observed fraction filled within the
// horizon? The gate blocks a challenger that calibrates materially worse.
type TimeToFillEval struct {
	Horizon   int
	Evaluated int
	FillRate  float64
	ChampPred float64
	ChalPred  float64
	ChampECE  float64
	ChalECE   float64
	Pass      bool
}

// EvaluateTimeToFill reads the pooled requirement durations, holds out a deterministic tail, and over the held
// out requirements whose fill within horizon outcome is KNOWN (filled at or before the horizon, or observed open
// for at least the horizon) compares the calibration of the champion and the challenger predicted fill within
// horizon probabilities against the observed fill rate. A requirement still open but observed for less than the
// horizon has an unknown outcome and is excluded. The challenger passes when its calibration error is within
// tolerance of the champion.
func EvaluateTimeToFill(ctx context.Context, pool *pgxpool.Pool, champ, chal Model, horizon int, holdout float64) (TimeToFillEval, error) {
	durations, err := ReadRequirementDurations(ctx, pool, HorizonMonths)
	if err != nil {
		return TimeToFillEval{}, err
	}
	if len(durations) == 0 {
		return TimeToFillEval{}, fmt.Errorf("magstore: no requirements to evaluate")
	}
	if horizon <= 0 {
		horizon = 3
	}
	if holdout <= 0 || holdout >= 1 {
		holdout = 0.30
	}
	cut := len(durations) - int(float64(len(durations))*holdout)
	if cut >= len(durations) {
		cut = len(durations) - 1
	}
	if cut < 0 {
		cut = 0
	}
	test := durations[cut:]

	res := TimeToFillEval{Horizon: horizon}
	res.ChampPred = FillWithinHorizon(champ, horizon)
	res.ChalPred = FillWithinHorizon(chal, horizon)

	var y []int
	for _, d := range test {
		known := (d.Filled && d.Periods <= horizon) || d.Periods >= horizon
		if !known {
			continue // still open, observed for less than the horizon: the outcome within the horizon is unknown
		}
		if d.Filled && d.Periods <= horizon {
			y = append(y, 1)
		} else {
			y = append(y, 0)
		}
	}
	res.Evaluated = len(y)
	if len(y) == 0 {
		res.Pass = true // nothing decidable on the held out split; do not block on an empty comparison
		return res, nil
	}
	var sum int
	for _, v := range y {
		sum += v
	}
	res.FillRate = float64(sum) / float64(len(y))
	champProbs := make([]float64, len(y))
	chalProbs := make([]float64, len(y))
	for i := range y {
		champProbs[i] = res.ChampPred
		chalProbs[i] = res.ChalPred
	}
	res.ChampECE = learn.ECE(champProbs, y, 1)
	res.ChalECE = learn.ECE(chalProbs, y, 1)
	res.Pass = res.ChalECE <= res.ChampECE+ECETolerance
	return res, nil
}