"""FastAPI server exposing SAM-family predictors over HTTP.

Implements the Label-Studio-ML `/predict` contract the SAM-Tool Tauri app
already speaks, and an LS-orthogonal `/models` endpoint so the frontend can
offer a model picker.

Run:
    python server.py --host 0.0.0.0 --port 9090 --device cuda

Endpoints:
    GET  /            → liveness + backend info
    GET  /models      → registered models + availability
    POST /predict     → LS-ML protocol, accepts optional `?model=<id>` query
"""

from __future__ import annotations

import argparse
import base64
import io
import logging
import os
import sys
from contextlib import asynccontextmanager
from pathlib import Path

import numpy as np
import uvicorn
from fastapi import FastAPI, HTTPException, Query
from fastapi.middleware.cors import CORSMiddleware
from PIL import Image
from pydantic import BaseModel

from registry import Registry, load_registry

# --------------------------------------------------------------------------
# Logging
# --------------------------------------------------------------------------
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(name)s: %(message)s",
    stream=sys.stdout,
)
log = logging.getLogger("sam2-backend.server")

# --------------------------------------------------------------------------
# App + registry. Registry is loaded inside a FastAPI lifespan so it works
# regardless of how uvicorn ends up importing this module — `python
# server.py` runs us as `__main__` and uvicorn then re-imports as `server`,
# so initialising at module top-level (or only in main()) leaves the live
# copy uninitialised.
# --------------------------------------------------------------------------
REGISTRY: Registry | None = None


@asynccontextmanager
async def lifespan(_app: FastAPI):
    global REGISTRY
    if REGISTRY is None:
        models_path = os.environ.get(
            "SAM_BACKEND_MODELS",
            str(Path(__file__).parent / "models.yaml"),
        )
        device = _pick_device(os.environ.get("SAM_BACKEND_DEVICE", "auto"))
        log.info("lifespan: loading registry (models=%s device=%s)",
                 models_path, device)
        REGISTRY = load_registry(Path(models_path), device=device)
        log.info(
            "lifespan: %d model specs (%d available), default=%s",
            len(REGISTRY.specs),
            sum(1 for s in REGISTRY.specs.values() if s.available),
            REGISTRY.default_id,
        )
    yield


app = FastAPI(title="SAM Backend", version="0.1.0", lifespan=lifespan)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


def _registry() -> Registry:
    if REGISTRY is None:
        raise RuntimeError("registry not initialised")
    return REGISTRY


# --------------------------------------------------------------------------
# Schemas
# --------------------------------------------------------------------------
class ModelInfo(BaseModel):
    id: str
    name: str
    family: str
    available: bool
    checkpoint: str


class ModelListResponse(BaseModel):
    default: str | None
    models: list[ModelInfo]


class HealthResponse(BaseModel):
    status: str
    device: str
    models: int
    default: str | None


# --------------------------------------------------------------------------
# Routes
# --------------------------------------------------------------------------
@app.get("/", response_model=HealthResponse)
def health() -> HealthResponse:
    reg = _registry()
    return HealthResponse(
        status="ok",
        device=reg.device,
        models=len(reg.specs),
        default=reg.default_id,
    )


@app.get("/models", response_model=ModelListResponse)
def list_models() -> ModelListResponse:
    reg = _registry()
    return ModelListResponse(
        default=reg.default_id,
        models=[
            ModelInfo(
                id=s.id,
                name=s.name,
                family=s.family,
                available=s.available,
                checkpoint=s.checkpoint_abs,
            )
            for s in reg.list_specs()
        ],
    )


@app.post("/predict")
def predict(payload: dict, model: str | None = Query(default=None)) -> dict:
    """Label-Studio-ML protocol.

    Body shape (simplified — we ignore fields we don't need):
        {
          "tasks": [{ "data": { "image": "data:image/jpeg;base64,..." } }],
          "params": {
            "context": {
              "result": [{
                "original_width": W, "original_height": H,
                "value": { "x": %, "y": %, "width": %, "height": %, ... }
              }]
            }
          }
        }

    Response shape (what sam-tool-tauri / LS Studio expects):
        {
          "results": [{
            "model_version": "<id>",
            "result": [{
              "original_width": W, "original_height": H,
              "value": {
                "points": [[x%, y%], ...],
                "closed": true,
                "polygonlabels": ["Object"]
              },
              "from_name": "polygon", "to_name": "image",
              "type": "polygonlabels"
            }]
          }]
        }
    """
    reg = _registry()

    try:
        tasks = payload.get("tasks") or []
        if not tasks:
            raise HTTPException(400, "tasks[] missing or empty")
        data = tasks[0].get("data") or {}
        image_field = data.get("image")
        if not isinstance(image_field, str):
            raise HTTPException(400, "tasks[0].data.image missing")
        image_rgb = _decode_image(image_field)

        ctx_result = (
            payload.get("params", {})
                   .get("context", {})
                   .get("result", [])
        )
        if not ctx_result:
            raise HTTPException(400, "params.context.result[] missing (bbox prompt)")
        prompt = ctx_result[0]
        value = prompt.get("value", {})
        orig_w = int(prompt.get("original_width") or image_rgb.shape[1])
        orig_h = int(prompt.get("original_height") or image_rgb.shape[0])
        x_pct = float(value.get("x", 0.0))
        y_pct = float(value.get("y", 0.0))
        w_pct = float(value.get("width", 0.0))
        h_pct = float(value.get("height", 0.0))

        # Percentages → absolute pixel xyxy, clamped to image bounds.
        x1 = max(0, int(round(x_pct / 100.0 * orig_w)))
        y1 = max(0, int(round(y_pct / 100.0 * orig_h)))
        x2 = min(orig_w, int(round((x_pct + w_pct) / 100.0 * orig_w)))
        y2 = min(orig_h, int(round((y_pct + h_pct) / 100.0 * orig_h)))
        if x2 <= x1 or y2 <= y1:
            raise HTTPException(400, "degenerate bbox after clamping")

        # Model selection: query param wins, then payload.params.model, then default.
        model_id = (
            model
            or payload.get("params", {}).get("model")
            or payload.get("params", {}).get("model_id")
        )
        # Resolve up-front so an unknown id is a clean 400; any error from
        # the actual predict path falls into the generic 500 with traceback
        # below (and is no longer mis-reported as "unknown model").
        try:
            reg.resolve(model_id)
        except KeyError as e:
            raise HTTPException(400, f"unknown model: {e}")

        spec, result = reg.predict_bbox(model_id, image_rgb, (x1, y1, x2, y2))

    except HTTPException:
        raise
    except FileNotFoundError as e:
        raise HTTPException(503, str(e))
    except Exception as e:
        log.exception("predict failed")
        raise HTTPException(500, f"{type(e).__name__}: {e}")

    if not result.polygon:
        raise HTTPException(500, "no polygon returned from model")

    # Convert pixel polygon back to percentages for LS-ML round-trip.
    points_pct = [
        [px / orig_w * 100.0, py / orig_h * 100.0]
        for (px, py) in result.polygon
    ]

    return {
        "results": [{
            "model_version": spec.id,
            "score": result.score,
            "result": [{
                "original_width": orig_w,
                "original_height": orig_h,
                "image_rotation": 0,
                "value": {
                    "points": points_pct,
                    "closed": True,
                    "polygonlabels": ["Object"],
                },
                "id": "sam_poly_0",
                "from_name": "polygon",
                "to_name": "image",
                "type": "polygonlabels",
            }],
        }]
    }


# --------------------------------------------------------------------------
# Helpers
# --------------------------------------------------------------------------
def _decode_image(field: str) -> np.ndarray:
    """Accept `data:<mime>;base64,...` or a bare base64 string."""
    b64 = field
    if field.startswith("data:"):
        comma = field.find(",")
        if comma == -1:
            raise HTTPException(400, "malformed data URL")
        b64 = field[comma + 1:]
    try:
        raw = base64.b64decode(b64)
    except Exception as e:
        raise HTTPException(400, f"base64 decode failed: {e}")
    try:
        img = Image.open(io.BytesIO(raw)).convert("RGB")
    except Exception as e:
        raise HTTPException(400, f"image decode failed: {e}")
    return np.array(img)


# --------------------------------------------------------------------------
# Entrypoint
# --------------------------------------------------------------------------
def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", default="0.0.0.0")
    parser.add_argument("--port", type=int, default=9090)
    parser.add_argument(
        "--models",
        default=str(Path(__file__).parent / "models.yaml"),
        help="Path to models.yaml registry file",
    )
    parser.add_argument(
        "--device",
        default="auto",
        help="'auto' (prefer cuda), 'cuda', 'cpu', or specific cuda:N",
    )
    parser.add_argument("--reload", action="store_true",
                        help="Enable uvicorn auto-reload (dev only)")
    args = parser.parse_args()

    # Hand config to the worker module via env — the lifespan reads these
    # whether uvicorn imports this file as `__main__` or as `server`.
    os.environ["SAM_BACKEND_MODELS"] = args.models
    os.environ["SAM_BACKEND_DEVICE"] = args.device

    log.info("starting uvicorn on %s:%d (device=%s, models=%s)",
             args.host, args.port, args.device, args.models)
    uvicorn.run(
        "server:app",
        host=args.host,
        port=args.port,
        reload=args.reload,
        log_level="info",
    )


def _pick_device(preference: str) -> str:
    if preference != "auto":
        return preference
    try:
        import torch
        if torch.cuda.is_available():
            return "cuda"
    except ImportError:
        pass
    return "cpu"


if __name__ == "__main__":
    main()