package engine

import (
	"predictor-refactored/internal/numerics"
)

// Propagator advances state under one Model, checking a set of Constraints
// after every integration step.
//
// When a constraint fires, the propagator binary-search refines the
// violation point and emits it as its final trajectory point. The Action of
// the triggering constraint controls what the surrounding Profile does
// next: stop the profile, transfer to Fallback, or clip and continue.
type Propagator struct {
	// Name identifies the propagator in trajectory metadata. Optional —
	// callers using sequential profile chains may leave it empty.
	Name string

	// Step is the magnitude of the integration step in seconds (always positive).
	// The Profile flips its sign for Reverse direction.
	Step float64

	// Model is the per-second derivative function used for integration.
	// One of Model or BuildModel must be non-nil. If both are set, BuildModel
	// takes precedence (it is invoked once per stage with a StageContext).
	Model      Model
	BuildModel func(ctx StageContext) Model

	// Constraints are evaluated after each step. The first violation wins.
	Constraints      []Constraint
	BuildConstraints func(ctx StageContext) []Constraint

	// Fallback is the propagator to switch to when a constraint with
	// ActionFallback fires. Optional.
	Fallback *Propagator

	// Tolerance is the binary-search refinement tolerance in parameter
	// space (default 0.01, matching Tawhiri).
	Tolerance float64
}

// run integrates the model from (t0, s0) in direction dir, returning a Result.
// globals are constraints injected by the Profile and checked alongside the
// propagator's local Constraints. events receives non-fatal observations.
func (p *Propagator) run(ctx StageContext, t0 float64, s0 State, globals []Constraint, events *EventSink) Result {
	dt := p.Step * float64(ctx.Direction)
	tol := p.Tolerance
	if tol == 0 {
		tol = 0.01
	}

	model := p.Model
	if p.BuildModel != nil {
		model = p.BuildModel(ctx)
	}
	constraints := p.Constraints
	if p.BuildConstraints != nil {
		constraints = p.BuildConstraints(ctx)
	}

	deriv := numerics.Deriv[State](func(t float64, s State) State { return model(t, s) })
	add := numerics.VecAdd[State](stateAdd)
	lerp := numerics.VecLerp[State](stateLerp)

	out := Result{
		Propagator: p.Name,
		Outcome:    OutcomeContinued,
		Points: []TrajectoryPoint{{
			Time: t0, Lat: s0.Lat, Lng: s0.Lng, Altitude: s0.Altitude,
		}},
	}

	t := t0
	s := s0

	for {
		s2 := numerics.RK4Step(t, s, dt, deriv, add)
		t2 := t + dt

		c, fired := firstFiring(constraints, globals, t2, s2)
		if !fired {
			t, s = t2, s2
			out.Points = append(out.Points, TrajectoryPoint{
				Time: t, Lat: s.Lat, Lng: s.Lng, Altitude: s.Altitude,
			})
			continue
		}

		// Record the unrefined violation.
		out.ViolationTime = t2
		out.ViolationState = s2

		trig := numerics.Trigger[State](func(tt float64, ss State) bool { return c.Violated(tt, ss) })
		t3, s3 := numerics.RefineTrigger(t, s, t2, s2, trig, lerp, tol)
		out.RefinedTime = t3
		out.RefinedState = s3
		out.Constraint = c
		out.ConstraintName = c.Name()

		switch c.Action() {
		case ActionClip:
			s3 = clipToConstraint(c, s3)
			out.RefinedState = s3
			out.Points = append(out.Points, TrajectoryPoint{
				Time: t3, Lat: s3.Lat, Lng: s3.Lng, Altitude: s3.Altitude,
			})
			t, s = t3, s3
			continue
		case ActionFallback:
			out.Points = append(out.Points, TrajectoryPoint{
				Time: t3, Lat: s3.Lat, Lng: s3.Lng, Altitude: s3.Altitude,
			})
			out.Outcome = OutcomeFallback
			out.Events = events.Snapshot()
			return out
		default: // ActionStop
			out.Points = append(out.Points, TrajectoryPoint{
				Time: t3, Lat: s3.Lat, Lng: s3.Lng, Altitude: s3.Altitude,
			})
			out.Outcome = OutcomeStopped
			out.Events = events.Snapshot()
			return out
		}
	}
}

// firstFiring scans local then global constraints for the first one whose
// Violated returns true at (t, s).
func firstFiring(local, globals []Constraint, t float64, s State) (Constraint, bool) {
	for _, c := range local {
		if c.Violated(t, s) {
			return c, true
		}
	}
	for _, c := range globals {
		if c.Violated(t, s) {
			return c, true
		}
	}
	return nil, false
}

// clipToConstraint adjusts s so that the given constraint is exactly
// satisfied (not violated). Defined only for constraints with a
// well-defined coordinate boundary; others fall through unchanged.
func clipToConstraint(c Constraint, s State) State {
	if alt, ok := c.(Altitude); ok {
		s.Altitude = alt.Limit
	}
	return s
}