// Command train-model is the offline TRAINER of the MAG learning loop. It fits an interpretable win
// probability model from the accruing bid_outcomes labels and writes a new CHALLENGER version of a MAG model
// (aedt_learned_models, aedt_model_coefficients, aedt_calibration_bins). It NEVER flips champion: promotion is
// a separate human governance action run through cmd/eval-model after the no adverse impact gate passes. The
// frozen decimal matching engine (internal/modules/matching, ModelVersion v1.4.0) stays the only per candidate
// scorer; this fits an OPPORTUNITY level win probability, never a candidate score.
//
// It reads the decided bids (win or loss) POOLED across every tenant. The MAG coefficient store is GLOBAL
// reference data (aedt_learned_models has no organization_id), so the learned model is necessarily one global,
// pooled model, not a per tenant model; a pooled training posture is the founder relevant decision the store's
// design already made, and cmd/eval-model's four fifths no adverse impact gate is what makes the pooled model
// safe to promote. The features are MARKET STRUCTURE only (the segment base rate), NONE a protected attribute
// or a proxy, and the SAME feature the serving path (internal/modules/ai gatherWinProbabilityFeatures) gathers,
// so a promoted challenger's learned coefficient lights up on the existing serve path with no serving change.
//
// The fit is iteratively reweighted least squares logistic regression in internal/platform/learn (pure Go, no
// Python), orchestrated by magstore.Train. float64 lives ONLY inside that offline solver; every coefficient
// this writes is re quantized to shopspring/decimal at the store's numeric scale, so the served path stays
// decimal exact.
//
// Env (server side only): MM_MIGRATOR_DATABASE_URL (or MM_DATABASE_URL), the BYPASSRLS migrator DSN.
// Flags: -model (the MAG model_key, default opportunity_win_probability), -l2 (ridge strength, default 1.0),
// -min-events (the events floor; 0 auto computes from the feature count), -dry-run (fit and report, write nothing).
//
// Build and run FIPS validated: CGO_ENABLED=0 GOFIPS140=v1.0.0 go build, run with GODEBUG=fips140=on.
package main

import (
	"context"
	"crypto/fips140"
	"crypto/pbkdf2"
	"crypto/sha256"
	"flag"
	"log"
	"os"
	"strings"
	"time"

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

	"github.com/missionworkforce/missionmatch/internal/platform/database"
	"github.com/missionworkforce/missionmatch/internal/platform/magstore"
)

func main() {
	kind := flag.String("kind", "win_probability", "which MAG model to fit: win_probability (opportunity_score) or reranker (display_order)")
	modelKey := flag.String("model", "", "the MAG model_key to fit a challenger for; defaults per kind")
	l2 := flag.Float64("l2", 1.0, "ridge strength (a zero mean Gaussian prior; a thin dataset shrinks toward the prior)")
	minEvents := flag.Int("min-events", 0, "the decided events (or pairs) floor; 0 auto computes from the feature count")
	dryRun := flag.Bool("dry-run", false, "fit and report the coefficients but write no challenger")
	flag.Parse()

	if err := fipsSelfTest(); err != nil {
		log.Fatalf("train-model: FIPS self test (PBKDF2 HMAC SHA 256) failed: %v", err)
	}
	log.Printf("train-model: fips_enabled=%v module=%s", fips140.Enabled(), fips140.Version())

	dsn := firstNonEmpty(os.Getenv("MM_MIGRATOR_DATABASE_URL"), os.Getenv("MM_DATABASE_URL"))
	if dsn == "" {
		log.Fatal("train-model: set MM_MIGRATOR_DATABASE_URL (the BYPASSRLS migrator DSN)")
	}

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute)
	defer cancel()

	pool, err := database.NewPool(ctx, dsn)
	if err != nil {
		log.Fatalf("train-model: connect: %v", err)
	}
	defer pool.Close()

	// The trainer reads labels pooled across tenants and writes the global coefficient store, both of which
	// require the BYPASSRLS migrator role. Fail closed if handed the app DSN.
	if err := database.AssertMigratorRole(ctx, pool); err != nil {
		log.Fatalf("train-model: migrator role check: %v", err)
	}

	if *kind == "reranker" {
		trainReranker(ctx, pool, resolveModelKey(*modelKey, magstore.RerankModelKey), *l2, *minEvents, *dryRun)
		return
	}
	if *kind == "timetofill" {
		trainTimeToFill(ctx, pool, resolveModelKey(*modelKey, magstore.TimeToFillModelKey), *minEvents, *dryRun)
		return
	}

	winModelKey := resolveModelKey(*modelKey, "opportunity_win_probability")
	res, err := magstore.Train(ctx, pool, winModelKey, *l2, *minEvents, *dryRun)
	if err != nil {
		log.Fatalf("train-model: %v", err)
	}
	log.Printf("train-model: %d decided bids (%d win, %d loss), %d features %v", res.N, res.Wins, res.Losses, len(res.FeatureNames), res.FeatureNames)
	if res.Skipped {
		log.Printf("train-model: %s; the prior only champion stays in force, nothing written", res.SkipReason)
		return
	}
	log.Printf("train-model: fit converged in %d iterations, l2=%v", res.Iterations, *l2)
	log.Printf("train-model:   intercept = %.10f", res.Intercept)
	for j, name := range res.FeatureNames {
		log.Printf("train-model:   %s = %.10f", name, res.Coeffs[j])
	}
	log.Printf("train-model: materialized %d calibration bins", res.Bins)
	if *dryRun {
		log.Print("train-model: dry run, no challenger written")
		return
	}
	log.Printf("train-model: wrote challenger %s v%d (champion=false); promote it with cmd/eval-model after the fairness gate passes", winModelKey, res.Version)
}

// trainReranker fits a pairwise learning to rank challenger of the display order re ranker over the decided
// slates and logs the learned utility weights. It NEVER flips champion; cmd/eval-model gates a promotion.
func trainReranker(ctx context.Context, pool *pgxpool.Pool, modelKey string, l2 float64, minPairs int, dryRun bool) {
	res, err := magstore.TrainReranker(ctx, pool, modelKey, l2, minPairs, dryRun)
	if err != nil {
		log.Fatalf("train-model: reranker: %v", err)
	}
	log.Printf("train-model: %d decided slates, %d approved over rejected pairs, %d features", res.Groups, res.Pairs, len(res.FeatureNames))
	if res.Skipped {
		log.Printf("train-model: %s; the prior only engine order champion stays in force, nothing written", res.SkipReason)
		return
	}
	log.Printf("train-model: fit pairwise utility weights (l2=%v)", l2)
	for j, name := range res.FeatureNames {
		log.Printf("train-model:   %s = %.10f", name, res.Coeffs[j])
	}
	if dryRun {
		log.Print("train-model: dry run, no challenger written")
		return
	}
	log.Printf("train-model: wrote reranker challenger %s v%d (champion=false); promote it with cmd/eval-model -kind reranker after the ranking and fairness gates pass", modelKey, res.Version)
}

// trainTimeToFill fits the empirical discrete time hazard challenger of the time to fill survival model over the
// pooled requirement durations and logs the per month hazard curve. It NEVER flips champion.
func trainTimeToFill(ctx context.Context, pool *pgxpool.Pool, modelKey string, minRequirements int, dryRun bool) {
	res, err := magstore.TrainTimeToFill(ctx, pool, modelKey, minRequirements, dryRun)
	if err != nil {
		log.Fatalf("train-model: timetofill: %v", err)
	}
	log.Printf("train-model: %d requirements (%d filled)", res.Requirements, res.Filled)
	if res.Skipped {
		log.Printf("train-model: %s; the prior only baseline hazard champion stays in force, nothing written", res.SkipReason)
		return
	}
	for m, h := range res.Hazard {
		log.Printf("train-model:   hazard_month_%d = %.6f", m+1, h)
	}
	if dryRun {
		log.Print("train-model: dry run, no challenger written")
		return
	}
	log.Printf("train-model: wrote time to fill challenger %s v%d (champion=false); promote it with cmd/eval-model -kind timetofill after the calibration gate passes", modelKey, res.Version)
}

// resolveModelKey returns an explicit model_key flag when set, else the default for the kind.
func resolveModelKey(flagKey, def string) string {
	if s := strings.TrimSpace(flagKey); s != "" {
		return s
	}
	return def
}

// fipsSelfTest derives one PBKDF2 HMAC SHA 256 key under the active crypto module and fails closed if it
// errors, so a build that dropped the validated module is caught at boot, the same gate the other cmds use.
func fipsSelfTest() error {
	_, err := pbkdf2.Key(sha256.New, "train-model-self-test", []byte("train-model-salt"), 600000, 32)
	return err
}

func firstNonEmpty(vs ...string) string {
	for _, v := range vs {
		if s := strings.TrimSpace(v); s != "" {
			return s
		}
	}
	return ""
}