app/main.py

import os
import pathlib
from datetime import datetime, timedelta, timezone
from typing import Dict, Any, List, Optional

import httpx
from fastapi import FastAPI, Query
from fastapi.responses import HTMLResponse, JSONResponse
from fastapi.staticfiles import StaticFiles

try:
    import eccodes as ec  # type: ignore
except Exception:
    ec = None


APP_TITLE = "ECMWF Open Data – 4‑Day Point Forecast"

app = FastAPI(title=APP_TITLE)


# Serve static files (index.html, JS, CSS)
app.mount("/static", StaticFiles(directory="app/static"), name="static")


def iso_utc(dt: datetime) -> str:
    return dt.astimezone(timezone.utc).strftime("%Y-%m-%dT%H:%M")


@app.get("/", response_class=HTMLResponse)
async def index():
    with open("app/static/index.html", "r", encoding="utf-8") as f:
        return HTMLResponse(f.read())


def map_vars() -> Dict[str, str]:
    """
    Map ECMWF short names to Open-Meteo variable names.
    Only include those available from Open-Meteo ECMWF IFS.
    Missing ones will be handled gracefully.
    """
    return {
        # 2t, 2d
        "2t": "temperature_2m",
        "2d": "dewpoint_2m",
        # 10u, 10v
        "10u": "wind_u_component_10m",
        "10v": "wind_v_component_10m",
        # msl, sp
        "msl": "pressure_msl",
        "sp": "surface_pressure",
        # total precipitation, convective precipitation, snowfall
        "tp": "precipitation",
        "cp": "convective_precipitation",  # may not always be available
        "sf": "snowfall",
        # cloud cover
        "tcc": "cloudcover",
        "lcc": "cloudcover_low",
        "mcc": "cloudcover_mid",
        "hcc": "cloudcover_high",
        # cape, cin, vis
        "cape": "cape",
        "cin": "convective_inhibition",  # may not always be available
        "vis": "visibility",
    }


def build_openmeteo_url(lat: float, lon: float, hours: int) -> str:
    """Build Open‑Meteo ECMWF URL for 4‑day hourly forecast.
    We request 4 days regardless of `hours`; the caller passes 96 for clarity.
    """
    base = "https://api.open-meteo.com/v1/ecmwf"
    vars_map = map_vars()
    hourly_vars = ",".join(sorted(set(v for v in vars_map.values())))

    params = {
        "latitude": f"{lat:.5f}",
        "longitude": f"{lon:.5f}",
        "hourly": hourly_vars,
        "forecast_days": "4",
        "windspeed_unit": "ms",
        "timezone": "UTC",
    }

    qp = httpx.QueryParams(params)
    return f"{base}?{qp}"


def translate_timeseries(api_json: Dict[str, Any]) -> Dict[str, Any]:
    """Translate Open-Meteo variable names back to ECMWF short names where possible."""
    if not api_json or "hourly" not in api_json:
        return {"hours": [], "data": {}}

    hours = api_json["hourly"].get("time", [])
    data: Dict[str, List[Any]] = {}

    inv_map = {v: k for k, v in map_vars().items()}
    for var_name, values in api_json["hourly"].items():
        if var_name == "time":
            continue
        short = inv_map.get(var_name, var_name)
        data[short] = values

    return {"hours": hours, "data": data, "source": "Open-Meteo ECMWF IFS"}


@app.get("/api/forecast")
async def forecast(lat: float = Query(...), lon: float = Query(...)):
    # Serve from locally cached ECMWF files only; do not download here.
    try:
        payload = await forecast_ecmwf(lat, lon)
        payload["source"] = "ECMWF Open Data (IFS HRES 0.4°)"
        return JSONResponse(payload)
    except RuntimeError as e:
        # Not ready yet; let frontend know cache is warming
        return JSONResponse(status_code=202, content={"warming": True, "message": str(e)})
    except Exception as e:
        # As a last resort, try Open‑Meteo
        url = build_openmeteo_url(lat, lon, hours=96)
        async with httpx.AsyncClient(timeout=60) as client:
            try:
                r = await client.get(url)
                r.raise_for_status()
                data = r.json()
                payload = translate_timeseries(data)
                payload["lat"] = lat
                payload["lon"] = lon
                return JSONResponse(payload)
            except httpx.HTTPError:
                return JSONResponse(status_code=502, content={"error": "Failed to fetch forecast from ECMWF and Open‑Meteo"})


# ---------------- ECMWF Open Data implementation ----------------

def _init_cache_dir() -> pathlib.Path:
    candidates = []
    env = os.getenv("ECMWF_CACHE")
    if env:
        candidates.append(pathlib.Path(env))
    # Prefer writable ephemeral storage in containers
    candidates += [
        pathlib.Path("/tmp/ecmwf-cache"),
        pathlib.Path("./cache"),
    ]
    for p in candidates:
        try:
            p.mkdir(parents=True, exist_ok=True)
            test = p / ".writetest"
            with open(test, "w") as f:
                f.write("ok")
            test.unlink(missing_ok=True)  # type: ignore[call-arg]
            return p
        except Exception:
            continue
    # Last resort: current directory (may still fail at runtime)
    return pathlib.Path(".")


CACHE_DIR = _init_cache_dir()

WANTED_SHORTNAMES = [
    "2t", "2d", "10u", "10v", "msl", "sp", "tp", "cp", "sf",
    "tcc", "lcc", "mcc", "hcc", "cape", "cin", "vis",
]


def latest_run(now: Optional[datetime] = None) -> tuple[str, str]:
    """Return (YYYYMMDD, HH) for latest reasonably available run (allowing latency)."""
    now = now or datetime.now(timezone.utc)
    # Allow 3h latency
    now -= timedelta(hours=3)
    hour = (now.hour // 6) * 6
    t = datetime(now.year, now.month, now.day, hour, tzinfo=timezone.utc)
    return t.strftime("%Y%m%d"), t.strftime("%H")


def step_list_4d() -> List[int]:
    # 6‑hourly steps up to 96 to keep size manageable
    return list(range(0, 97, 6))


def sfc_grib_filename(date: str, hour: str, step: int) -> str:
    return f"fc_sfc_0p4-{date}_{hour}00-{step:03d}h.grib2"


def sfc_grib_url(date: str, hour: str, step: int) -> str:
    # Construct known Open Data path for HRES sfc global file
    # Path format: /forecasts/YYYYMMDD/HHz/ifs/0p4/oper/fc_sfc_0p4-YYYYMMDD_HH00-XXXh.grib2
    return (
        f"https://data.ecmwf.int/forecasts/{date}/{hour}z/ifs/0p4/oper/"
        f"{sfc_grib_filename(date, hour, step)}"
    )


async def http_download(url: str, target: pathlib.Path, timeout: int = 120) -> None:
    if target.exists() and target.stat().st_size > 0:
        return
    target.parent.mkdir(parents=True, exist_ok=True)
    async with httpx.AsyncClient(timeout=timeout, follow_redirects=True) as client:
        r = await client.get(url)
        r.raise_for_status()
        with open(target, "wb") as f:
            f.write(r.content)


def extract_point_from_grib(path: pathlib.Path, lat: float, lon: float, wanted: List[str], step: Optional[int] = None) -> Dict[str, float]:
    vals: Dict[str, float] = {}
    if ec is None:
        return vals
    with open(path, "rb") as f:
        while True:
            gid = ec.codes_grib_new_from_file(f)
            if gid is None:
                break
            try:
                short = ec.codes_get(gid, "shortName")
                if short in wanted and short not in vals:
                    if step is not None:
                        try:
                            st = int(ec.codes_get(gid, "step"))
                        except Exception:
                            st = None
                        if st != step:
                            continue
                    nearest = ec.codes_grib_find_nearest(gid, lat, lon)[0]
                    vals[short] = float(nearest["value"])  # type: ignore[index]
                if len(vals) == len(wanted):
                    break
            finally:
                ec.codes_release(gid)
    return vals


def clamp(v: float, lo: float, hi: float) -> float:
    return max(lo, min(hi, v))


async def forecast_ecmwf(lat: float, lon: float) -> Dict[str, Any]:
    if ec is None:
        raise RuntimeError("ecCodes not available")

    date, hour = latest_run()
    steps = step_list_4d()

    # Use only locally cached files
    hours_iso: List[str] = []
    data: Dict[str, List[Optional[float]]] = {k: [] for k in WANTED_SHORTNAMES}
    run_base = datetime.strptime(f"{date} {hour}", "%Y%m%d %H").replace(tzinfo=timezone.utc)

    available = 0
    for s in steps:
        local = CACHE_DIR / date / hour / sfc_grib_filename(date, hour, s)
        if not local.exists() or local.stat().st_size == 0:
            continue
        available += 1
        vals = extract_point_from_grib(local, lat, lon, WANTED_SHORTNAMES)
        ts = run_base + timedelta(hours=s)
        hours_iso.append(ts.strftime("%Y-%m-%dT%H:%M"))
        for k in WANTED_SHORTNAMES:
            data[k].append(vals.get(k))

    if available == 0:
        raise RuntimeError("ECMWF cache warming; no steps available yet")

    return {"lat": lat, "lon": lon, "hours": hours_iso, "data": data}


# Background prefetcher to download ECMWF GRIBs on startup
import asyncio

prefetch_state: Dict[str, Any] = {"run": None, "total": 0, "ready": 0, "bytes": 0}


async def prefetch_latest() -> None:
    date, hour = latest_run()
    steps = step_list_4d()
    prefetch_state["run"] = {"date": date, "hour": hour}
    prefetch_state["total"] = len(steps)
    prefetch_state["ready"] = 0
    prefetch_state["bytes"] = 0
    for s in steps:
        try:
            local = CACHE_DIR / date / hour / sfc_grib_filename(date, hour, s)
            url = sfc_grib_url(date, hour, s)
            await http_download(url, local)
            prefetch_state["ready"] = sum(1 for st in steps if (CACHE_DIR / date / hour / sfc_grib_filename(date, hour, st)).exists())
            prefetch_state["bytes"] = sum((CACHE_DIR / date / hour / sfc_grib_filename(date, hour, st)).stat().st_size for st in steps if (CACHE_DIR / date / hour / sfc_grib_filename(date, hour, st)).exists())
        except Exception:
            continue


@app.on_event("startup")
async def _startup():
    asyncio.create_task(prefetch_latest())


@app.get("/api/status")
async def status():
    return {
        "cache_dir": str(CACHE_DIR),
        "run": prefetch_state.get("run"),
        "total": prefetch_state.get("total"),
        "ready": prefetch_state.get("ready"),
        "bytes": prefetch_state.get("bytes"),
    }


def get_port() -> int:
    try:
        return int(os.getenv("PORT", "7860"))
    except Exception:
        return 7860


if __name__ == "__main__":
    import uvicorn

    uvicorn.run("app.main:app", host="0.0.0.0", port=get_port(), reload=False)