predictor/internal/api/prediction.go

package api

import (
	"context"
	"fmt"
	"net/http"
	"time"

	"predictor-refactored/internal/engine"
	"predictor-refactored/internal/weather"
	apirest "predictor-refactored/pkg/rest"
)

// ReadinessCheck implements GET /ready.
func (h *Handler) ReadinessCheck(_ context.Context) (*apirest.ReadinessResponse, error) {
	resp := &apirest.ReadinessResponse{}
	if field := h.mgr.Active(); field != nil {
		resp.Status = apirest.ReadinessResponseStatusOk
		resp.DatasetTime = apirest.NewOptDateTime(field.Epoch())
	} else {
		resp.Status = apirest.ReadinessResponseStatusNotReady
		resp.ErrorMessage = apirest.NewOptString("no dataset loaded")
	}
	return resp, nil
}

// fieldFor resolves the dataset a request asked for, or the active one.
//
// A request naming a dataset that is not stored is refused. The parameter used
// to be declared in the spec and never read, so a client asking for a specific
// forecast run — an archival date, a reproducible comparison — was silently
// served from whatever happened to be loaded instead.
func (h *Handler) fieldFor(ctx context.Context, want apirest.OptDateTime) (weather.WindField, error) {
	if epoch, ok := want.Get(); ok {
		field, err := h.mgr.FieldFor(ctx, epoch.UTC())
		if err != nil {
			return nil, apiError(http.StatusBadRequest, err.Error())
		}
		return field, nil
	}
	field := h.mgr.Active()
	if field == nil {
		return nil, apiError(http.StatusServiceUnavailable, "no dataset loaded, service is starting up")
	}
	return field, nil
}

// runFailure turns an integration failure into a 400.
//
// The sampler reporting an error means it had no data — most often a launch time
// past the dataset's horizon — so the trajectory is not a forecast. It used to be
// discarded, and the run continued with zero wind: the balloon took off and
// landed on the same spot, returned as a successful prediction.
func runFailure(events *engine.EventSink, field weather.WindField) error {
	err := events.Err()
	if err == nil {
		return nil
	}
	return apiError(http.StatusBadRequest,
		fmt.Sprintf("%s (dataset %s)", err, field.Epoch().UTC().Format(time.RFC3339)))
}

// PerformPredictionV2 implements POST /api/v2/prediction.
func (h *Handler) PerformPredictionV2(ctx context.Context, req *apirest.PredictionV2Request) (*apirest.PredictionV2Response, error) {
	resp, err := h.runPredictionV2(ctx, req)
	if err == nil {
		h.metrics.Prediction("v2", resp.CompletedAt.Sub(resp.StartedAt), nil)
	}
	return resp, err
}

// CreatePredictionJob implements POST /api/v1/predictions.
func (h *Handler) CreatePredictionJob(_ context.Context, req *apirest.PredictionV2Request) (*apirest.PredictionJob, error) {
	info, accepted := h.async.Enqueue(req)
	if !accepted {
		return nil, apiError(http.StatusServiceUnavailable, info.Error)
	}
	return asyncJobToAPI(info), nil
}

// GetPredictionJob implements GET /api/v1/predictions/{id}.
func (h *Handler) GetPredictionJob(_ context.Context, params apirest.GetPredictionJobParams) (*apirest.PredictionJob, error) {
	info, ok := h.async.Get(params.ID)
	if !ok {
		return nil, apiError(http.StatusNotFound, "prediction job not found")
	}
	return asyncJobToAPI(info), nil
}

// CancelPredictionJob implements DELETE /api/v1/predictions/{id}.
func (h *Handler) CancelPredictionJob(_ context.Context, params apirest.CancelPredictionJobParams) error {
	if !h.async.Cancel(params.ID) {
		return apiError(http.StatusConflict, "job not found or already terminal")
	}
	return nil
}

// runPredictionV2 is the synchronous prediction core, shared by the v2
// endpoint and the async worker pool.
func (h *Handler) runPredictionV2(ctx context.Context, req *apirest.PredictionV2Request) (*apirest.PredictionV2Response, error) {
	// Validate the request shape before checking dataset availability, so a
	// malformed request is a 400 regardless of startup state.
	lat := req.Launch.Latitude
	alt := req.Launch.Altitude.Or(0)
	if err := validateLat(lat); err != nil {
		return nil, err
	}
	if err := validateLng(req.Launch.Longitude); err != nil {
		return nil, err
	}
	lng := normalizeLng(req.Launch.Longitude)

	field, err := h.fieldFor(ctx, req.Dataset)
	if err != nil {
		return nil, err
	}

	events := engine.NewEventSink()
	deps := engine.BuildDeps{Wind: field, Events: events, Terrain: h.terrain()}

	prof, err := buildProfile(req, deps)
	if err != nil {
		return nil, apiError(http.StatusBadRequest, err.Error())
	}

	started := time.Now().UTC()
	results := prof.Run(float64(req.Launch.Time.Unix()), engine.State{Lat: lat, Lng: lng, Altitude: alt}, events)
	completed := time.Now().UTC()
	if err := runFailure(events, field); err != nil {
		return nil, err
	}

	resp := &apirest.PredictionV2Response{
		Stages:      make([]apirest.StageResult, 0, len(results)),
		Events:      eventsToAPI(events.Snapshot()),
		Dataset:     apirest.DatasetInfo{Source: field.Source(), Epoch: field.Epoch()},
		StartedAt:   started,
		CompletedAt: completed,
	}
	for _, r := range results {
		resp.Stages = append(resp.Stages, stageResultToAPI(r))
	}
	return resp, nil
}

// PerformPrediction implements GET /api/v1/prediction (Tawhiri-compatible).
func (h *Handler) PerformPrediction(ctx context.Context, params apirest.PerformPredictionParams) (*apirest.PredictionResponse, error) {
	field, err := h.fieldFor(ctx, params.Dataset)
	if err != nil {
		return nil, err
	}

	profileKind := "standard_profile"
	if p, ok := params.Profile.Get(); ok {
		profileKind = string(p)
	}
	ascentRate := params.AscentRate.Or(5)
	descentRate := params.DescentRate.Or(5)
	launchAlt := params.LaunchAltitude.Or(0)
	if err := validateLat(params.LaunchLatitude); err != nil {
		return nil, err
	}
	if err := validateLng(params.LaunchLongitude); err != nil {
		return nil, err
	}
	lng := normalizeLng(params.LaunchLongitude)
	launchTime := float64(params.LaunchDatetime.Unix())

	events := engine.NewEventSink()
	var stageNames []string
	var prof engine.Profile
	switch profileKind {
	case "standard_profile":
		stageNames = []string{"ascent", "descent"}
		prof = standardProfile(field, h.terrain(), events, ascentRate, params.BurstAltitude.Or(28000), descentRate)
	case "float_profile":
		stopTime := params.LaunchDatetime.Add(24 * time.Hour)
		if v, ok := params.StopDatetime.Get(); ok {
			stopTime = v
		}
		stageNames = []string{"ascent", "float"}
		prof = floatProfile(field, events, ascentRate, params.FloatAltitude.Or(25000), stopTime)
	default:
		return nil, apiError(http.StatusBadRequest, "unknown profile: "+profileKind)
	}

	started := time.Now().UTC()
	results := prof.Run(launchTime, engine.State{Lat: params.LaunchLatitude, Lng: lng, Altitude: launchAlt}, events)
	completed := time.Now().UTC()
	if err := runFailure(events, field); err != nil {
		return nil, err
	}
	h.metrics.Prediction(profileKind, completed.Sub(started), nil)

	resp := &apirest.PredictionResponse{
		Metadata: apirest.PredictionResponseMetadata{StartDatetime: started, CompleteDatetime: completed},
	}
	for i, r := range results {
		name := "ascent"
		if i < len(stageNames) {
			name = stageNames[i]
		}
		resp.Prediction = append(resp.Prediction, tawhiriItem(name, r))
	}
	resp.Request = apirest.NewOptPredictionResponseRequest(apirest.PredictionResponseRequest{
		Dataset:         apirest.NewOptString(field.Epoch().Format("2006-01-02T15:04:05Z")),
		LaunchLatitude:  apirest.NewOptFloat64(params.LaunchLatitude),
		LaunchLongitude: apirest.NewOptFloat64(params.LaunchLongitude),
		LaunchDatetime:  apirest.NewOptString(params.LaunchDatetime.Format(time.RFC3339)),
		LaunchAltitude:  params.LaunchAltitude,
	})
	if ev := events.Snapshot(); len(ev) > 0 {
		resp.Warnings = apirest.NewOptPredictionResponseWarnings(apirest.PredictionResponseWarnings{})
	}
	return resp, nil
}

// standardProfile builds the Tawhiri ascent → descent chain.
func standardProfile(field weather.WindField, elev engine.TerrainProvider, events *engine.EventSink, ascentRate, burst, descentRate float64) engine.Profile {
	wind := engine.WindTransport(field, events)
	descentTerm := []engine.Constraint{engine.Altitude{Op: engine.OpLessEqual, Limit: 0, On: engine.ActionStop}}
	if elev != nil {
		descentTerm = []engine.Constraint{engine.TerrainContact{Provider: elev, On: engine.ActionStop}}
	}
	return engine.Profile{
		Direction: engine.Forward,
		Stages: []*engine.Propagator{
			{
				Name:        "ascent",
				Step:        60,
				Model:       engine.Sum(engine.ConstantRate(ascentRate), wind),
				Constraints: []engine.Constraint{engine.Altitude{Op: engine.OpGreaterEqual, Limit: burst, On: engine.ActionStop}},
			},
			{
				Name:        "descent",
				Step:        60,
				Model:       engine.Sum(engine.ParachuteDescent(descentRate), wind),
				Constraints: descentTerm,
			},
		},
	}
}

// floatProfile builds the Tawhiri ascent → float chain.
func floatProfile(field weather.WindField, events *engine.EventSink, ascentRate, floatAlt float64, stopTime time.Time) engine.Profile {
	wind := engine.WindTransport(field, events)
	return engine.Profile{
		Direction: engine.Forward,
		Stages: []*engine.Propagator{
			{
				Name:        "ascent",
				Step:        60,
				Model:       engine.Sum(engine.ConstantRate(ascentRate), wind),
				Constraints: []engine.Constraint{engine.Altitude{Op: engine.OpGreaterEqual, Limit: floatAlt, On: engine.ActionStop}},
			},
			{
				Name:        "float",
				Step:        60,
				Model:       wind,
				Constraints: []engine.Constraint{engine.Time{Op: engine.OpGreater, Limit: float64(stopTime.Unix()), On: engine.ActionStop}},
			},
		},
	}
}

// tawhiriItem maps one engine stage result to a v1 prediction item.
func tawhiriItem(name string, r engine.Result) apirest.PredictionResponsePredictionItem {
	stage := apirest.PredictionResponsePredictionItemStageAscent
	switch name {
	case "descent":
		stage = apirest.PredictionResponsePredictionItemStageDescent
	case "float":
		stage = apirest.PredictionResponsePredictionItemStageFloat
	}
	n := r.Path.Len()
	traj := make([]apirest.TawhiriPoint, 0, n)
	for i := range n {
		t, p := r.Path.At(i)
		traj = append(traj, apirest.TawhiriPoint{
			Datetime: time.Unix(int64(t), 0).UTC(),
			Latitude: p.Lat,
			// v1 publishes [0, 360), the convention upstream Tawhiri publishes —
			// measured: it answers 308.7884 where we used to answer -51.2115 for
			// the same meridian. v1 exists to be a drop-in, so it matches. v2 keeps
			// the signed format (mapping.go).
			Longitude: p.Lng,
			Altitude:  p.Altitude,
		})
	}
	return apirest.PredictionResponsePredictionItem{Stage: stage, Trajectory: traj}
}