"""
FastAPI JSON API for the bundled EWS model.
Run locally:  uvicorn main:app --host 0.0.0.0 --port 7860
"""

from __future__ import annotations

import os
from pathlib import Path

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field

import ews_inference_core as ews

BUNDLE_PATH = Path(os.environ.get("EWS_BUNDLE_PATH", "/app/ews_pipeline_bundle.joblib"))
if not BUNDLE_PATH.is_file():
    BUNDLE_PATH = Path(__file__).resolve().parent / "ews_pipeline_bundle.joblib"

_bundle: dict | None = None

app = FastAPI(title="MP Dengue EWS Inference", version="1.0")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.on_event("startup")
def _load() -> None:
    global _bundle
    p = BUNDLE_PATH
    if not p.is_file():
        raise RuntimeError(
            f"Bundle not found: {p}. Run bundle_ews_pipeline.py and copy the .joblib here, "
            "or set EWS_BUNDLE_PATH."
        )
    _bundle = ews.load_bundle(str(p))


def _bundle_get() -> dict:
    if _bundle is None:
        raise HTTPException(status_code=503, detail="Model not loaded")
    return _bundle


class PredictRequest(BaseModel):
    """Provide either `features` (name → value) or `values` (length = n_features in order)."""

    features: dict[str, float] | None = None
    values: list[float] | None = None
    population: int | None = Field(
        default=None,
        ge=1,
        description="District population for estimated_cases and risk_tier (optional)",
    )


@app.get("/")
def root():
    from fastapi.responses import RedirectResponse

    return RedirectResponse(url="/docs")


@app.get("/health")
def health():
    return {"status": "ok", "bundle": str(BUNDLE_PATH.resolve())}


@app.get("/schema")
def schema():
    b = _bundle_get()
    cfg = b["config"]
    return {
        "feature_columns": cfg["feature_columns"],
        "n_features": cfg["n_features"],
        "prediction_horizon_weeks": cfg.get("prediction_horizon"),
        "cases_lags": cfg.get("cases_lags"),
        "probability_threshold": cfg.get("decision_threshold"),
    }


@app.post("/predict")
def predict(body: PredictRequest):
    b = _bundle_get()
    if body.features is not None and body.values is not None:
        raise HTTPException(
            status_code=400,
            detail="Send only one of `features` or `values`, not both.",
        )
    if body.features is None and body.values is None:
        raise HTTPException(
            status_code=400,
            detail="Provide `features` (object) or `values` (array).",
        )
    try:
        if body.features is not None:
            X = ews.vector_from_feature_dict(b, body.features)
        else:
            X = ews.vector_from_values(b, body.values or [])
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e)) from e

    try:
        out = ews.predict_from_bundle(b, X, body.population)
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Prediction failed: {e!s}") from e
    return {"ok": True, "result": out}