#!/usr/bin/env python3
"""Precompute oriented detections for demo images via the Modal API.

Calls the structured ``object-predict`` endpoint (bboxes + labels + scores)
at a low score threshold so the static UI can filter client-side.

Reads DL4EO_API_KEY from the environment or a repo-root ``.env`` file.
"""

from __future__ import annotations

import argparse
import base64
import json
import os
import sys
import time
from pathlib import Path

import numpy as np
import requests
from PIL import Image

ROOT = Path(__file__).resolve().parents[1]
DEMO_DIR = ROOT / "demo"
OUT_DIR = ROOT / "predictions"
ENV_FILE = ROOT / ".env"

DEFAULT_API_URL = "https://dl4eo--object-predict.modal.run"
REQUEST_TIMEOUT_SEC = 180
MAX_RETRIES = 3
RETRYABLE_STATUS_CODES = {408, 429, 502, 503, 504}


def load_dotenv(path: Path = ENV_FILE) -> None:
    """Load KEY=VALUE pairs from .env into os.environ (does not override existing)."""
    if not path.is_file():
        return
    for raw in path.read_text().splitlines():
        line = raw.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        key, _, value = line.partition("=")
        key = key.strip()
        value = value.strip().strip("'").strip('"')
        if key and key not in os.environ:
            os.environ[key] = value


def predict(image: np.ndarray, threshold: float, api_url: str, api_key: str) -> dict:
    payload = {
        "image": base64.b64encode(np.ascontiguousarray(image)).decode(),
        "shape": list(image.shape),
        "threshold": threshold,
    }
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }

    last_error = None
    response = None
    for attempt in range(1, MAX_RETRIES + 1):
        try:
            response = requests.post(
                api_url,
                json=payload,
                headers=headers,
                timeout=REQUEST_TIMEOUT_SEC,
            )
            if response.status_code == 200:
                return response.json()
            if response.status_code not in RETRYABLE_STATUS_CODES:
                raise RuntimeError(
                    f"API status={response.status_code}: {response.text[:300]}"
                )
            last_error = f"status={response.status_code}"
        except requests.exceptions.RequestException as exc:
            last_error = str(exc)
            response = None
        if attempt < MAX_RETRIES:
            time.sleep(2**attempt)

    raise RuntimeError(f"API failed after {MAX_RETRIES} attempts ({last_error})")


def process_image(
    image_path: Path,
    threshold: float,
    api_url: str,
    api_key: str,
) -> dict:
    pil = Image.open(image_path).convert("RGB")
    arr = np.asarray(pil)
    raw = predict(arr, threshold=threshold, api_url=api_url, api_key=api_key)

    detections = []
    for bbox, label, score in zip(raw["bboxes"], raw["labels"], raw["scores"]):
        detections.append(
            {
                "bbox": [float(x) for x in bbox],
                "label": str(label),
                "score": float(score),
            }
        )
    detections.sort(key=lambda d: d["score"], reverse=True)

    return {
        "image": f"demo/{image_path.name}",
        "width": int(arr.shape[1]),
        "height": int(arr.shape[0]),
        "threshold_used": float(threshold),
        "duration": float(raw.get("duration", 0.0)),
        "detections": detections,
    }


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--threshold",
        type=float,
        default=0.05,
        help="Minimum score kept when calling Modal (default: 0.05)",
    )
    parser.add_argument(
        "--api-url",
        default=os.environ.get("DL4EO_API_URL", DEFAULT_API_URL),
        help="Modal object-predict endpoint URL",
    )
    parser.add_argument(
        "--image",
        action="append",
        dest="images",
        help="Optional demo filename (repeatable). Default: all demo/*.jpg",
    )
    args = parser.parse_args()

    load_dotenv()
    api_key = os.environ.get("DL4EO_API_KEY")
    if not api_key:
        print(
            f"ERROR: set DL4EO_API_KEY in the environment or in {ENV_FILE}",
            file=sys.stderr,
        )
        return 1

    if args.images:
        paths = [DEMO_DIR / name for name in args.images]
    else:
        paths = sorted(DEMO_DIR.glob("*.jpg"))

    if not paths:
        print("ERROR: no demo images found", file=sys.stderr)
        return 1

    OUT_DIR.mkdir(parents=True, exist_ok=True)
    manifest = {"demos": []}

    for path in paths:
        if not path.exists():
            print(f"ERROR: missing {path}", file=sys.stderr)
            return 1
        print(f"Predicting {path.name} …")
        result = process_image(
            path,
            threshold=args.threshold,
            api_url=args.api_url,
            api_key=api_key,
        )
        out_path = OUT_DIR / f"{path.stem}.json"
        out_path.write_text(json.dumps(result, indent=2) + "\n")
        print(
            f"  → {out_path.relative_to(ROOT)} "
            f"({len(result['detections'])} dets, {result['duration']:.2f}s)"
        )
        manifest["demos"].append(
            {
                "id": path.stem,
                "title": path.stem.replace("_", " "),
                "image": result["image"],
                "predictions": f"predictions/{path.stem}.json",
                "width": result["width"],
                "height": result["height"],
                "count": len(result["detections"]),
            }
        )

    manifest_path = OUT_DIR / "manifest.json"
    manifest_path.write_text(json.dumps(manifest, indent=2) + "\n")
    print(f"Wrote {manifest_path.relative_to(ROOT)}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())