package magstore

import (
	"context"
	"fmt"
	"sort"
	"strings"

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

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

// RankTolerance is the slack the ranking quality gate allows so a tiny held out set or a floating point tie does
// not fail a genuinely equal challenger. The challenger passes when its pairwise accuracy is within this of the
// champion (better or equal).
const RankTolerance = 0.02

// RerankModelKey is the model_key of the learned slate display order re ranker (migration 0037 seeds the prior
// only v1 whose coefficients are the frozen engine component weights, so it reproduces the engine order until a
// challenger is fitted and promoted).
const RerankModelKey = "slate_display_order"

// RerankFeatureOrder is the FIXED feature order the display order re ranker trains and serves on, the frozen
// engine's six sub scores. Training, ordering, and the coefficient store all use this exact order so a fitted
// coefficient pairs with its sub score. NONE is a protected attribute or a proxy.
var RerankFeatureOrder = []string{"clearance_score", "skills_score", "certs_score", "customer_score", "experience_score", "location_score"}

// SlateMember is one decided slate member for the re ranker: its slate and candidate ids, its six frozen engine
// sub score features (aligned to RerankFeatureOrder), and whether the human approved it. The ids carry NO
// demographic data; they are the only thing that crosses to the eeo compliance path for the fairness gate.
type SlateMember struct {
	SlateID     string
	CandidateID string
	Features    []float64
	Approved    bool
}

// SlateGroup is one slate's decided members, the ranking group the pairwise re ranker learns within. OrgID
// carries the owning tenant so the fairness gate can recompute the four fifths ratio per org (a tenant scoped
// EEO measure) under a challenger ranking.
type SlateGroup struct {
	SlateID string
	OrgID   string
	Members []SlateMember
}

// ReadSlateGroups reads every decided slate member (approved or rejected) POOLED across tenants, grouped by
// slate, with the six frozen engine sub scores as float features. It reads ONLY the frozen sub scores and the
// decision, never a demographic attribute (the eeo schema is unreachable from here); the demographic join for
// the fairness gate happens separately through the compliance pool keyed only by candidate id. Ordered
// deterministically by slate then candidate.
func ReadSlateGroups(ctx context.Context, pool *pgxpool.Pool) ([]SlateGroup, error) {
	const q = `
SELECT slate_id::text, organization_id::text, candidate_id::text, decision,
       clearance_score::float8, skills_score::float8, certs_score::float8,
       customer_score::float8, experience_score::float8, location_score::float8
  FROM slate_members
 WHERE decision IN ('approved', 'rejected')
 ORDER BY slate_id, candidate_id`
	rows, err := pool.Query(ctx, q)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var groups []SlateGroup
	var cur *SlateGroup
	for rows.Next() {
		var slateID, orgID, candidateID, decision string
		f := make([]float64, len(RerankFeatureOrder))
		if err := rows.Scan(&slateID, &orgID, &candidateID, &decision, &f[0], &f[1], &f[2], &f[3], &f[4], &f[5]); err != nil {
			return nil, err
		}
		if cur == nil || cur.SlateID != slateID {
			groups = append(groups, SlateGroup{SlateID: slateID, OrgID: orgID})
			cur = &groups[len(groups)-1]
		}
		cur.Members = append(cur.Members, SlateMember{
			SlateID:     slateID,
			CandidateID: candidateID,
			Features:    f,
			Approved:    decision == "approved",
		})
	}
	return groups, rows.Err()
}

// toRankGroups projects the slate groups into the pure learn ranking groups (features plus the approved label),
// dropping the ids the numerical core does not need.
func toRankGroups(groups []SlateGroup) []learn.RankGroup {
	out := make([]learn.RankGroup, 0, len(groups))
	for _, g := range groups {
		items := make([]learn.RankItem, 0, len(g.Members))
		for _, m := range g.Members {
			items = append(items, learn.RankItem{Features: m.Features, Approved: m.Approved})
		}
		out = append(out, learn.RankGroup{Items: items})
	}
	return out
}

// RerankTrainResult summarizes a re ranker training run.
type RerankTrainResult struct {
	Skipped      bool
	SkipReason   string
	Version      int
	Groups       int
	Pairs        int
	FeatureNames []string
	Coeffs       []float64 // raw utility weights, aligned to RerankFeatureOrder
}

// TrainReranker reads the decided slates, fits a pairwise learning to rank model over the frozen sub score
// features (approved should rank above rejected within a slate), and writes a CHALLENGER version of the display
// order re ranker. It refuses to fit below the pair floor, leaving the prior only engine order champion in
// force. It NEVER flips champion. The learned utility weights order candidates; a display order model emits an
// ORDERING, never a per candidate score, and stores no calibration bins.
func TrainReranker(ctx context.Context, pool *pgxpool.Pool, modelKey string, l2 float64, minPairs int, dryRun bool) (RerankTrainResult, error) {
	groups, err := ReadSlateGroups(ctx, pool)
	if err != nil {
		return RerankTrainResult{}, err
	}
	x, y := learn.PairwiseDataset(toRankGroups(groups))
	res := RerankTrainResult{Groups: len(groups), Pairs: len(x), FeatureNames: RerankFeatureOrder}

	if minPairs <= 0 {
		minPairs = 20 * len(RerankFeatureOrder) // twenty pairs per feature, the ranking analogue of the events floor
	}
	if len(x) < minPairs {
		res.Skipped = true
		res.SkipReason = fmt.Sprintf("%d approved over rejected pairs is below the floor of %d", len(x), minPairs)
		return res, nil
	}

	model, err := learn.FitLogistic(x, y, l2)
	if err != nil {
		return RerankTrainResult{}, err
	}
	// The utility weights: the raw space coefficients. The intercept is irrelevant to an ordering (it is a
	// constant across candidates), so it is not stored as a weight; the display order is the utility descending.
	res.Coeffs = make([]float64, len(RerankFeatureOrder))
	for j := range RerankFeatureOrder {
		res.Coeffs[j] = model.Coeffs[j] / model.Scales[j]
	}

	if dryRun {
		return res, nil
	}

	notes := fmt.Sprintf("Challenger pairwise learning to rank fit by cmd/train-model over %d decided slates (%d approved over rejected pairs) pooled across tenants, ridge l2=%v. OutputKind display_order, never a candidate score; the frozen v1.4.0 engine stays the only per candidate scorer and a human approves every member. champion=false until cmd/eval-model passes the ranking and four fifths gates and a human promotes it.", len(groups), len(x), l2)
	version, err := WriteChallenger(ctx, pool, modelKey, Challenger{
		Family:       "pairwise_logistic",
		OutputKind:   "display_order",
		FeatureNames: RerankFeatureOrder,
		Intercept:    0, // an ordering is invariant to the intercept
		Coeffs:       res.Coeffs,
		Notes:        notes,
	})
	if err != nil {
		return RerankTrainResult{}, err
	}
	res.Version = version
	return res, nil
}

// RerankCoeffs extracts a loaded model's utility weights in RerankFeatureOrder, so a stored re ranker (champion
// or challenger) can order a slate or be scored for ranking quality.
func RerankCoeffs(m Model) []float64 {
	out := make([]float64, len(RerankFeatureOrder))
	for j, name := range RerankFeatureOrder {
		out[j] = m.Coeffs[name]
	}
	return out
}

// OrderMembers returns the slate members in the display order the utility weights imply: the learned utility
// descending, with the frozen engine order (the incoming slice order) as the STABLE tiebreak, so a zero weight
// or a tie preserves the engine order exactly. It emits an ORDERING (a permutation), never a score.
func OrderMembers(coeffs []float64, members []SlateMember) []SlateMember {
	out := make([]SlateMember, len(members))
	copy(out, members)
	sort.SliceStable(out, func(i, j int) bool {
		return learn.RankUtility(out[i].Features, coeffs) > learn.RankUtility(out[j].Features, coeffs)
	})
	return out
}

// RerankQualityResult is the held out ranking comparison of the champion and a challenger re ranker.
type RerankQualityResult struct {
	HeldOutGroups int
	ChampAccuracy float64
	ChalAccuracy  float64
	Pass          bool
}

// EvaluateRerankerQuality reads the decided slates, holds out a deterministic tail fraction as a test split, and
// computes the pairwise ranking accuracy (approved ordered above rejected within a slate) under both the
// champion and the challenger utility weights. The challenger passes when it ranks at least as well as the
// champion (accuracy within RankTolerance). It is pure over the stored weights and the slate read; the fairness
// gate is a separate check the caller runs through the compliance pool.
func EvaluateRerankerQuality(ctx context.Context, pool *pgxpool.Pool, champCoeffs, chalCoeffs []float64, holdout float64) (RerankQualityResult, error) {
	groups, err := ReadSlateGroups(ctx, pool)
	if err != nil {
		return RerankQualityResult{}, err
	}
	if len(groups) == 0 {
		return RerankQualityResult{}, fmt.Errorf("magstore: no decided slates to evaluate")
	}
	if holdout <= 0 || holdout >= 1 {
		holdout = 0.30
	}
	cut := len(groups) - int(float64(len(groups))*holdout)
	if cut >= len(groups) {
		cut = len(groups) - 1
	}
	if cut < 0 {
		cut = 0
	}
	test := toRankGroups(groups[cut:])
	res := RerankQualityResult{
		HeldOutGroups: len(groups) - cut,
		ChampAccuracy: learn.PairwiseAccuracy(test, champCoeffs),
		ChalAccuracy:  learn.PairwiseAccuracy(test, chalCoeffs),
	}
	res.Pass = res.ChalAccuracy >= res.ChampAccuracy-RankTolerance
	return res, nil
}

// GroupsByOrg buckets the slate groups by owning tenant, so the fairness gate can recompute the four fifths
// ratio per org (a tenant scoped EEO measure) under a challenger ranking.
func GroupsByOrg(groups []SlateGroup) map[string][]SlateGroup {
	out := map[string][]SlateGroup{}
	for _, g := range groups {
		out[g.OrgID] = append(out[g.OrgID], g)
	}
	return out
}

// CandidateSelection reduces the per slate top k selection to a per candidate set (selected if the ranking
// surfaces the candidate in the top of ANY of their slates), the shape the eeo four fifths recomputation
// consumes. It carries only candidate ids, no demographic data.
func CandidateSelection(coeffs []float64, groups []SlateGroup) map[string]bool {
	perSlate := TopKSelection(coeffs, groups)
	out := map[string]bool{}
	for key := range perSlate {
		if i := strings.LastIndex(key, "|"); i >= 0 {
			out[key[i+1:]] = true
		}
	}
	return out
}

// TopKSelection marks, per slate, the candidates the ranking would surface as the SELECTION it nudges: the top
// N by the utility weights where N is the number the human actually approved on that slate. It returns a set
// keyed by "slateID|candidateID". This is the selection proxy the fairness gate recomputes the four fifths
// ratio over, so a display order change that would skew who surfaces first is caught. Ties preserve the engine
// order (the stable incoming member order).
func TopKSelection(coeffs []float64, groups []SlateGroup) map[string]bool {
	selected := map[string]bool{}
	for _, g := range groups {
		n := 0
		for _, m := range g.Members {
			if m.Approved {
				n++
			}
		}
		if n == 0 {
			continue
		}
		ordered := OrderMembers(coeffs, g.Members)
		for i := 0; i < n && i < len(ordered); i++ {
			selected[ordered[i].SlateID+"|"+ordered[i].CandidateID] = true
		}
	}
	return selected
}