package evolve

import (
	"context"
	"errors"
	"fmt"
	"math"
	"sort"
)

// FeatureFunc maps a Candidate to a single scalar feature value. Pure; any
// error should be treated as a programmer bug (the caller picked a feature
// that cannot be evaluated on this candidate shape). If your feature
// genuinely can fail at runtime, route through FuncEvaluator instead.
type FeatureFunc func(Candidate) float64

// WeightedEvaluator is the reference "cheap scorer" for the common weighted-
// linear fitness pattern: fitness = sum_i (feature_i(candidate) * weight_i).
//
// Why this is the main reference impl, not a closure wrapper:
// Weighted-linear scoring covers the vast majority of "cheap first-pass"
// evaluators across domains (logistics carrier scoring, ad eCPM, multi-factor
// stock selection, hiring rubrics, engineering trade-offs). The shape is
// stable, declarative, and auditable; consumers plug in features without
// writing a full Evaluator.
//
// For non-linear / fn-style scoring use FuncEvaluator.
type WeightedEvaluator struct {
	features  map[string]FeatureFunc
	weights   map[string]float64
	normalize bool
}

// weightedConfig is the option-builder accumulator. Separated so option
// errors surface from NewWeightedEvaluator instead of being deferred to
// first Score call.
type weightedConfig struct {
	features  map[string]FeatureFunc
	weights   map[string]float64
	normalize bool
	// errs accumulates option-time errors so users see the full config
	// issue set in one NewWeightedEvaluator call, not one-per-retry.
	errs []error
}

// WeightedOption configures a WeightedEvaluator.
type WeightedOption func(*weightedConfig)

// WithFeature registers a named feature extractor. Duplicate names produce
// a constructor error (silently overwriting a feature is a frequent config-
// drift source: a copy-pasted option block overriding a prior definition).
func WithFeature(name string, fn FeatureFunc) WeightedOption {
	return func(c *weightedConfig) {
		if name == "" {
			c.errs = append(c.errs, errors.New("evolve: WithFeature requires non-empty name"))
			return
		}
		if fn == nil {
			c.errs = append(c.errs, fmt.Errorf("evolve: WithFeature(%q) requires non-nil fn", name))
			return
		}
		if _, dup := c.features[name]; dup {
			c.errs = append(c.errs, fmt.Errorf("evolve: WithFeature(%q) duplicate registration", name))
			return
		}
		c.features[name] = fn
	}
}

// WithWeight sets a single weight. Negative values are allowed and convey
// "smaller feature value is better" (e.g. cost, latency, error rate).
func WithWeight(name string, w float64) WeightedOption {
	return func(c *weightedConfig) {
		if name == "" {
			c.errs = append(c.errs, errors.New("evolve: WithWeight requires non-empty name"))
			return
		}
		if math.IsNaN(w) || math.IsInf(w, 0) {
			c.errs = append(c.errs, fmt.Errorf("evolve: WithWeight(%q) requires finite weight, got %v", name, w))
			return
		}
		c.weights[name] = w
	}
}

// WithWeights sets weights in bulk. Equivalent to WithWeight per entry.
func WithWeights(m map[string]float64) WeightedOption {
	return func(c *weightedConfig) {
		for name, w := range m {
			if name == "" {
				c.errs = append(c.errs, errors.New("evolve: WithWeights: empty name in map"))
				continue
			}
			if math.IsNaN(w) || math.IsInf(w, 0) {
				c.errs = append(c.errs, fmt.Errorf("evolve: WithWeights(%q) requires finite weight, got %v", name, w))
				continue
			}
			c.weights[name] = w
		}
	}
}

// WithNormalization enables [0, 1] clipping of every feature value before the
// weighted sum. Default off: many real-world features (cost in RMB, latency
// in ms, inventory count) are absolute scales where clipping silently corrupts
// the score. Only enable when your features are already bounded-unit
// (fraction, score 0-1).
func WithNormalization(on bool) WeightedOption {
	return func(c *weightedConfig) { c.normalize = on }
}

// NewWeightedEvaluator builds the evaluator and validates:
//   - at least one feature registered
//   - every feature has exactly one weight
//   - every weight has exactly one feature (no orphan weights)
//
// All errors are reported together (joined via errors.Join) so a misconfigured
// evaluator surfaces every issue in one go.
func NewWeightedEvaluator(opts ...WeightedOption) (*WeightedEvaluator, error) {
	cfg := weightedConfig{
		features: make(map[string]FeatureFunc),
		weights:  make(map[string]float64),
	}
	for _, o := range opts {
		o(&cfg)
	}

	var errs []error
	errs = append(errs, cfg.errs...)

	if len(cfg.features) == 0 {
		errs = append(errs, errors.New("evolve: WeightedEvaluator requires at least one feature"))
	}

	// Every feature must have a weight; otherwise its contribution is silently
	// zero, which is a lurking bug.
	for name := range cfg.features {
		if _, ok := cfg.weights[name]; !ok {
			errs = append(errs, fmt.Errorf("evolve: feature %q has no weight", name))
		}
	}
	// Every weight must refer to a known feature; orphan weights are user
	// error (typo in name, removed feature).
	for name := range cfg.weights {
		if _, ok := cfg.features[name]; !ok {
			errs = append(errs, fmt.Errorf("evolve: weight %q has no matching feature", name))
		}
	}

	if err := errors.Join(errs...); err != nil {
		return nil, err
	}

	return &WeightedEvaluator{
		features:  cfg.features,
		weights:   cfg.weights,
		normalize: cfg.normalize,
	}, nil
}

// Score implements Evaluator. breakdown contains the RAW feature values, not
// weighted contributions: a debugger wants to see "on_time feature was 0.3",
// not "contribution 0.15 = 0.3 * 0.5". Reconstructing contributions from
// breakdown + weights is a one-liner; the reverse is lossy.
func (e *WeightedEvaluator) Score(ctx context.Context, c Candidate) (float64, map[string]float64, error) {
	if err := ctx.Err(); err != nil {
		return 0, nil, err
	}
	breakdown := make(map[string]float64, len(e.features))
	var total float64
	// Iterate in sorted name order so the weighted sum is deterministic
	// regardless of map iteration order (important for reproducible fitness
	// values across runs when float accumulation order matters).
	names := make([]string, 0, len(e.features))
	for name := range e.features {
		names = append(names, name)
	}
	sort.Strings(names)
	for _, name := range names {
		v := e.features[name](c)
		if e.normalize {
			v = math.Max(0, math.Min(1, v))
		}
		breakdown[name] = v
		total += v * e.weights[name]
	}
	return total, breakdown, nil
}

// Features returns the registered feature names in sorted order. Useful for
// UI introspection and tests. Returns a copy; mutating the result does not
// affect the evaluator.
func (e *WeightedEvaluator) Features() []string {
	out := make([]string, 0, len(e.features))
	for name := range e.features {
		out = append(out, name)
	}
	sort.Strings(out)
	return out
}

// Weights returns a copy of the weight map.
func (e *WeightedEvaluator) Weights() map[string]float64 {
	out := make(map[string]float64, len(e.weights))
	for k, v := range e.weights {
		out[k] = v
	}
	return out
}

// ============================================================================
// FuncEvaluator: escape hatch for non-linear / bespoke scoring
// ============================================================================

// FuncEvaluator wraps a user-supplied scoring function as an Evaluator.
//
// Use this when your scorer is not well expressed as a weighted sum: e.g.
// a decision tree, a hard-constraint-first gate, a composed multi-stage
// score, or a function that genuinely needs to return an error (reading a
// lookup file, calling a side-channel service). Everything else should be
// a WeightedEvaluator -- it is more auditable and its coefficients can be
// evolved via ParameterStore.
type FuncEvaluator struct {
	fn func(ctx context.Context, c Candidate) (float64, map[string]float64, error)
}

// NewFuncEvaluator wraps fn. fn must be non-nil.
func NewFuncEvaluator(fn func(ctx context.Context, c Candidate) (float64, map[string]float64, error)) (*FuncEvaluator, error) {
	if fn == nil {
		return nil, errors.New("evolve: NewFuncEvaluator requires non-nil fn")
	}
	return &FuncEvaluator{fn: fn}, nil
}

// Score implements Evaluator by delegating to the wrapped function. ctx
// cancellation is checked before delegation; the wrapped fn is responsible
// for checking ctx itself on longer operations.
func (e *FuncEvaluator) Score(ctx context.Context, c Candidate) (float64, map[string]float64, error) {
	if err := ctx.Err(); err != nil {
		return 0, nil, err
	}
	return e.fn(ctx, c)
}