package engine

import (
	"fmt"
	"sync"

	"predictor-refactored/internal/weather"
)

// ConstraintSpec is the source-agnostic JSON-shape used to declare a
// constraint. The Type field is the registry key; remaining fields are
// extracted by the registered factory.
type ConstraintSpec struct {
	Type   string `json:"type"`
	Action string `json:"action,omitempty"`
	// Op is the comparison operator for scalar constraints (altitude, time).
	Op    string  `json:"op,omitempty"`
	Limit float64 `json:"limit,omitempty"`
	// Vertices and Mode are used by the polygon constraint.
	Vertices []PolygonVertex `json:"vertices,omitempty"`
	Mode     string          `json:"mode,omitempty"`
	// Label is an optional human-readable identifier surfaced via Name().
	Label string `json:"label,omitempty"`
}

// ModelSpec is the source-agnostic JSON shape used to declare a model.
type ModelSpec struct {
	Type string `json:"type"`
	// Rate (m/s) for constant_rate.
	Rate float64 `json:"rate,omitempty"`
	// SeaLevelRate (m/s, positive) for parachute_descent.
	SeaLevelRate float64 `json:"sea_level_rate,omitempty"`
	// Segments for piecewise.
	Segments []PiecewiseSegmentSpec `json:"segments,omitempty"`
	// IncludeWind sums a WindTransport model into the resulting derivative.
	IncludeWind bool `json:"include_wind,omitempty"`
}

// PiecewiseSegmentSpec is one entry in a piecewise rate schedule.
//
// Reference selects how the Until field is interpreted:
//
//   - "absolute" (default): UNIX seconds.
//   - "profile_start": seconds since the profile's launch time.
//   - "propagator_start": seconds since this propagator began running.
type PiecewiseSegmentSpec struct {
	Until     float64 `json:"until"`
	Rate      float64 `json:"rate"`
	Reference string  `json:"reference,omitempty"`
}

// BuildDeps bundle the runtime dependencies factories may consult.
type BuildDeps struct {
	Wind    weather.WindField
	Terrain TerrainProvider
	Events  *EventSink
}

// ConstraintFactory builds one Constraint from a spec.
type ConstraintFactory func(spec ConstraintSpec, deps BuildDeps) (Constraint, error)

// ModelFactory builds one model from a spec. The returned Built is held by
// a Propagator; if Build is set, it is invoked lazily by the profile
// runner before every stage so it can capture per-stage start times.
type ModelFactory func(spec ModelSpec, deps BuildDeps) (BuiltModel, error)

// BuiltModel is either an eager Model, a lazy Build, or both. The profile
// runner prefers Build when present.
type BuiltModel struct {
	Model Model
	Build func(ctx StageContext) Model
}

var (
	regMu               sync.RWMutex
	constraintFactories = map[string]ConstraintFactory{}
	modelFactories      = map[string]ModelFactory{}
)

// RegisterConstraint installs a factory for typeName. Subsequent calls
// overwrite the previous factory.
func RegisterConstraint(typeName string, f ConstraintFactory) {
	regMu.Lock()
	defer regMu.Unlock()
	constraintFactories[typeName] = f
}

// RegisterModel installs a model factory.
func RegisterModel(typeName string, f ModelFactory) {
	regMu.Lock()
	defer regMu.Unlock()
	modelFactories[typeName] = f
}

// BuildConstraint dispatches spec to its registered factory.
func BuildConstraint(spec ConstraintSpec, deps BuildDeps) (Constraint, error) {
	regMu.RLock()
	f, ok := constraintFactories[spec.Type]
	regMu.RUnlock()
	if !ok {
		return nil, fmt.Errorf("unknown constraint type %q", spec.Type)
	}
	return f(spec, deps)
}

// BuildModel dispatches spec to its registered factory.
func BuildModel(spec ModelSpec, deps BuildDeps) (BuiltModel, error) {
	regMu.RLock()
	f, ok := modelFactories[spec.Type]
	regMu.RUnlock()
	if !ok {
		return BuiltModel{}, fmt.Errorf("unknown model type %q", spec.Type)
	}
	return f(spec, deps)
}

// RegisteredConstraints returns the names of every registered constraint type.
func RegisteredConstraints() []string {
	regMu.RLock()
	defer regMu.RUnlock()
	out := make([]string, 0, len(constraintFactories))
	for k := range constraintFactories {
		out = append(out, k)
	}
	return out
}

// RegisteredModels returns the names of every registered model type.
func RegisteredModels() []string {
	regMu.RLock()
	defer regMu.RUnlock()
	out := make([]string, 0, len(modelFactories))
	for k := range modelFactories {
		out = append(out, k)
	}
	return out
}

// --- Built-in registrations ------------------------------------------------

func init() {
	RegisterConstraint("altitude", buildAltitude)
	RegisterConstraint("time", buildTime)
	RegisterConstraint("terrain_contact", buildTerrainContact)
	RegisterConstraint("polygon", buildPolygon)

	RegisterModel("constant_rate", buildConstantRate)
	RegisterModel("parachute_descent", buildParachuteDescent)
	RegisterModel("piecewise", buildPiecewise)
	RegisterModel("wind", buildWind)
}

func buildAltitude(spec ConstraintSpec, _ BuildDeps) (Constraint, error) {
	op, err := ParseOperator(spec.Op)
	if err != nil {
		return nil, fmt.Errorf("altitude: %w", err)
	}
	act, err := ParseAction(spec.Action)
	if err != nil {
		return nil, fmt.Errorf("altitude: %w", err)
	}
	return Altitude{Op: op, Limit: spec.Limit, On: act}, nil
}

func buildTime(spec ConstraintSpec, _ BuildDeps) (Constraint, error) {
	op, err := ParseOperator(spec.Op)
	if err != nil {
		return nil, fmt.Errorf("time: %w", err)
	}
	act, err := ParseAction(spec.Action)
	if err != nil {
		return nil, fmt.Errorf("time: %w", err)
	}
	return Time{Op: op, Limit: spec.Limit, On: act}, nil
}

func buildTerrainContact(spec ConstraintSpec, deps BuildDeps) (Constraint, error) {
	if deps.Terrain == nil {
		return nil, fmt.Errorf("terrain_contact requires a terrain provider")
	}
	act, err := ParseAction(spec.Action)
	if err != nil {
		return nil, fmt.Errorf("terrain_contact: %w", err)
	}
	return TerrainContact{Provider: deps.Terrain, On: act}, nil
}

func buildPolygon(spec ConstraintSpec, _ BuildDeps) (Constraint, error) {
	if len(spec.Vertices) < 3 {
		return nil, fmt.Errorf("polygon requires at least 3 vertices")
	}
	act, err := ParseAction(spec.Action)
	if err != nil {
		return nil, fmt.Errorf("polygon: %w", err)
	}
	mode := PolygonInside
	switch spec.Mode {
	case "", "inside":
		mode = PolygonInside
	case "outside":
		mode = PolygonOutside
	default:
		return nil, fmt.Errorf("polygon: unknown mode %q", spec.Mode)
	}
	return NewPolygon(spec.Vertices, mode, act, spec.Label), nil
}

func buildConstantRate(spec ModelSpec, _ BuildDeps) (BuiltModel, error) {
	return BuiltModel{Model: ConstantRate(spec.Rate)}, nil
}

func buildParachuteDescent(spec ModelSpec, _ BuildDeps) (BuiltModel, error) {
	if spec.SeaLevelRate <= 0 {
		return BuiltModel{}, fmt.Errorf("parachute_descent requires positive sea_level_rate")
	}
	return BuiltModel{Model: ParachuteDescent(spec.SeaLevelRate)}, nil
}

func buildWind(_ ModelSpec, deps BuildDeps) (BuiltModel, error) {
	if deps.Wind == nil {
		return BuiltModel{}, fmt.Errorf("wind model requires a loaded wind field")
	}
	return BuiltModel{Model: WindTransport(deps.Wind, deps.Events)}, nil
}

func buildPiecewise(spec ModelSpec, deps BuildDeps) (BuiltModel, error) {
	for _, s := range spec.Segments {
		switch s.Reference {
		case "", "absolute", "profile_start", "propagator_start":
		default:
			return BuiltModel{}, fmt.Errorf("piecewise: unknown segment reference %q", s.Reference)
		}
	}
	// Always build lazily: the profile runner supplies a StageContext before
	// each stage, which is what resolves absolute / profile-relative /
	// propagator-relative segment times uniformly.
	return BuiltModel{
		Build: func(ctx StageContext) Model {
			return maybeAddWind(Piecewise(resolveSegments(spec.Segments, ctx)), spec.IncludeWind, deps)
		},
	}, nil
}

// resolveSegments converts spec segments to engine.RateSegment, turning each
// segment's reference-relative Until into an absolute UNIX time. References
// are validated by buildPiecewise, so an unrecognised one here is treated as
// absolute rather than re-erroring.
func resolveSegments(in []PiecewiseSegmentSpec, ctx StageContext) []RateSegment {
	out := make([]RateSegment, 0, len(in))
	for _, s := range in {
		out = append(out, RateSegment{Until: segmentBase(s.Reference, ctx) + s.Until, Rate: s.Rate})
	}
	return out
}

// segmentBase returns the absolute time a piecewise segment's Until is
// measured from, per its reference.
func segmentBase(reference string, ctx StageContext) float64 {
	switch reference {
	case "profile_start":
		return ctx.ProfileStart
	case "propagator_start":
		return ctx.PropagatorStart
	default: // "", "absolute"
		return 0
	}
}

// maybeAddWind sums a WindTransport model into base when the spec asks for it.
func maybeAddWind(base Model, includeWind bool, deps BuildDeps) Model {
	if !includeWind {
		return base
	}
	if deps.Wind == nil {
		return base
	}
	return Sum(base, WindTransport(deps.Wind, deps.Events))
}