from __future__ import annotations

import io
import json
import os
import subprocess
import sys
import tempfile
import zipfile
from pathlib import Path

import numpy as np
import torch
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles
from PIL import Image, ImageDraw, ImageFont
from pydantic import BaseModel, Field

from gazelle.model_factory import get_gazelle_model
from gazelle.utils import visualize_heatmap


ROOT = Path(__file__).resolve().parent
STATIC_DIR = ROOT / "static"
CHECKPOINT_DIR = ROOT / "checkpoints"
BACKBONE_DIR = ROOT / "backbones"


def _checkpoint(env_var: str, default_name: str) -> Path:
    return Path(os.environ.get(env_var, str(CHECKPOINT_DIR / default_name)))


ALL_MODEL_OPTIONS = {
    "cross_gaze_vith_inout": {
        "source": "local_checkpoint",
        "checkpoint": _checkpoint("GAZELLE_VITH_CHECKPOINT", "vith.pt"),
        "factory_name": "cross_gaze_vith_inout",
        "include_backbone": True,
        "display_name": "CrossGaze ViT-H",
    },
    "cross_gaze_vitb_distill_screen": {
        "source": "local_checkpoint",
        "checkpoint": _checkpoint("GAZELLE_VITB_CHECKPOINT", "vitb_distill_screen.pt"),
        "factory_name": "cross_gaze_vitb_inout_finetune",
        "include_backbone": True,
        "display_name": "CrossGaze ViT-B Distill Screen",
    },
    "page_vitb": {
        "source": "hf_page",
        "repo_id": "Octopus1/page-vitb",
        "display_name": "PaGE ViT-B",
    },
    "page_vitb_screen": {
        "source": "hf_page",
        "repo_id": "Octopus1/page-vitb-screen",
        "display_name": "PaGE ViT-B Screen",
    },
    "page_vitsplus": {
        "source": "hf_page",
        "repo_id": "Octopus1/page-vitsplus",
        "display_name": "PaGE ViT-S+",
    },
    "page_vits": {
        "source": "hf_page",
        "repo_id": "Octopus1/page-vits",
        "display_name": "PaGE ViT-S",
    },
    "page_vithplus": {
        "source": "hf_page",
        "repo_id": "Octopus1/page-vithplus",
        "display_name": "PaGE ViT-H+",
    },
}

HF_ONLY_MODE = os.environ.get("SPACE_DEFAULT_HF_ONLY", "").lower() in {"1", "true", "yes"} or bool(os.environ.get("SPACE_ID"))
DEFAULT_MODEL_NAME = os.environ.get("DEFAULT_MODEL_NAME", "page_vithplus" if HF_ONLY_MODE else "cross_gaze_vitb_distill_screen")
PRELOAD_DEFAULT_MODEL = os.environ.get("PRELOAD_DEFAULT_MODEL", "0" if HF_ONLY_MODE else "1").lower() in {"1", "true", "yes"}
OUT_THRESHOLD = 0.5
MAX_INFER_SIDE = 720

MODEL_OPTIONS = {
    name: cfg
    for name, cfg in ALL_MODEL_OPTIONS.items()
    if not HF_ONLY_MODE or cfg.get("source") == "hf_page"
}

MODEL_PATH_ENV = {
    "cross_gaze_vith_inout": {
        "scene": "GAZELLE_DINOV3_VITH_SCENE_PATH",
        "head": "GAZELLE_DINOV3_VITH_HEAD_PATH",
    },
    "cross_gaze_vitb_inout_finetune": {
        "scene": "GAZELLE_DINOV3_VITB_SCENE_PATH",
        "head": "GAZELLE_DINOV3_VITB_HEAD_PATH",
    },
}

DEFAULT_FACTORY_PATHS = {
    "cross_gaze_vith_inout": {
        "scene": str(BACKBONE_DIR / "dinov3_vith16plus"),
        "head": str(BACKBONE_DIR / "dinov3_vith16plus"),
    },
    "cross_gaze_vitb_inout_finetune": {
        "scene": str(BACKBONE_DIR / "dinov3_vitb16"),
        "head": str(BACKBONE_DIR / "dinov3_vitb16"),
    },
}

HF_TOKEN_ENV_VARS = ("HUGGING_FACE_HUB_TOKEN", "HF_TOKEN")


def _resolve_backbone_paths(model_name: str) -> dict[str, str]:
    env_names = MODEL_PATH_ENV[model_name]
    defaults = DEFAULT_FACTORY_PATHS[model_name]
    resolved = {
        "scene": os.environ.get(env_names["scene"], defaults["scene"]),
        "head": os.environ.get(env_names["head"], defaults["head"]),
    }
    return resolved


def _patch_factory_backbone_paths(model_name: str):
    import gazelle.model_factory as model_factory_module

    resolved = _resolve_backbone_paths(model_name)
    original_scene = model_factory_module.DinoV3HFBackbone

    class PatchedDinoV3HFBackbone(original_scene):
        def __init__(self, model_path: str, in_size=(512, 512)):
            target_path = resolved["scene"] if tuple(in_size) == (512, 512) else resolved["head"]
            super().__init__(model_path=target_path, in_size=in_size)

    return model_factory_module, original_scene, PatchedDinoV3HFBackbone, resolved


class ModelManager:
    def __init__(self) -> None:
        self._cache: dict[str, tuple[torch.nn.Module, list, list | None]] = {}
        self.device = self._current_device()

    def _current_device(self) -> torch.device:
        return torch.device("cuda" if torch.cuda.is_available() else "cpu")

    def refresh_device(self) -> torch.device:
        next_device = self._current_device()
        if next_device != self.device:
            self.device = next_device
            for model, _, _ in self._cache.values():
                model.to(self.device)
        return self.device

    def move_to_cpu(self) -> None:
        self.device = torch.device("cpu")
        for model, _, _ in self._cache.values():
            model.to(self.device)
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    def get(self, model_name: str):
        self.refresh_device()
        cfg = MODEL_OPTIONS.get(model_name)
        if cfg is None:
            raise HTTPException(status_code=400, detail=f"Unsupported model: {model_name}")

        if model_name not in self._cache:
            if cfg.get("source") == "hf_page":
                model, scene_transforms, head_transforms = self._load_hf_page_model(cfg)
            else:
                model, scene_transforms, head_transforms = self._load_local_checkpoint_model(cfg)
            model.eval()
            model.to(self.device)
            self._cache[model_name] = (model, scene_transforms, head_transforms)

        return self._cache[model_name]

    def _load_local_checkpoint_model(self, cfg: dict):
        model_factory_module, original_scene, patched_scene, resolved_paths = _patch_factory_backbone_paths(cfg["factory_name"])
        model_factory_module.DinoV3HFBackbone = patched_scene
        try:
            model, scene_transforms, head_transforms = get_gazelle_model(cfg["factory_name"])
        except Exception as exc:
            detail = (
                f"Failed to build model. scene_path={resolved_paths['scene']} head_path={resolved_paths['head']} ; "
                f"set env vars {MODEL_PATH_ENV[cfg['factory_name']]['scene']} / {MODEL_PATH_ENV[cfg['factory_name']]['head']} if needed. "
                f"Original error: {exc}"
            )
            raise HTTPException(status_code=500, detail=detail) from exc
        finally:
            model_factory_module.DinoV3HFBackbone = original_scene

        checkpoint_path = cfg["checkpoint"]
        if not checkpoint_path.exists():
            raise HTTPException(status_code=500, detail=f"Checkpoint not found: {checkpoint_path}")

        checkpoint = self._load_checkpoint(checkpoint_path)
        model.load_gazelle_state_dict(checkpoint, include_backbone=cfg["include_backbone"])
        return model, scene_transforms, head_transforms

    def _load_hf_page_model(self, cfg: dict):
        from torchvision import transforms
        from transformers import AutoModel

        repo_id = cfg["repo_id"]
        token = next((os.environ.get(name) for name in HF_TOKEN_ENV_VARS if os.environ.get(name)), None)
        try:
            model = AutoModel.from_pretrained(
                repo_id,
                token=token,
                trust_remote_code=True,
            )
        except Exception as exc:
            raise HTTPException(
                status_code=500,
                detail=(
                    f"Failed to load Hugging Face model {repo_id}. "
                    f"For private repos, set HUGGING_FACE_HUB_TOKEN or HF_TOKEN. Original error: {exc}"
                ),
            ) from exc

        config = model.config
        image_mean = list(getattr(config, "image_mean", [0.485, 0.456, 0.406]))
        image_std = list(getattr(config, "image_std", [0.229, 0.224, 0.225]))
        scene_size = tuple(getattr(config, "scene_in_size", (512, 512)))
        head_size = tuple(getattr(config, "head_in_size", (256, 256)))

        def build_transform(size):
            return transforms.Compose([
                transforms.Resize(tuple(size)),
                transforms.ToTensor(),
                transforms.Normalize(mean=image_mean, std=image_std),
            ])

        return model, [build_transform(scene_size)], [build_transform(head_size)]

    def _load_checkpoint(self, checkpoint_path: Path):
        if zipfile.is_zipfile(checkpoint_path):
            with checkpoint_path.open("rb") as fh:
                if fh.read(4) != b"PK\x03\x04":
                    raise HTTPException(
                        status_code=500,
                        detail=(
                            f"Checkpoint {checkpoint_path} appears corrupted or padded: it has a ZIP central directory "
                            "but does not start with the PyTorch ZIP header. Please re-download or re-save the checkpoint."
                        ),
                    )
        try:
            checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
        except KeyError as exc:
            if "storages" not in str(exc):
                raise
            checkpoint = self._load_checkpoint_in_py310(checkpoint_path)

        if isinstance(checkpoint, dict) and "state_dict" in checkpoint:
            checkpoint = checkpoint["state_dict"]
        return self._normalize_checkpoint_keys(checkpoint)

    def _normalize_checkpoint_keys(self, checkpoint: dict[str, torch.Tensor]):
        normalized = {}
        replacements = {
            "scene_branch_backbone.backbone.model.model.": "scene_branch_backbone.backbone.model.",
            "head_branch_backbone.backbone.model.model.": "head_branch_backbone.backbone.model.",
        }
        for key, value in checkpoint.items():
            new_key = key
            for old, new in replacements.items():
                if new_key.startswith(old):
                    new_key = new + new_key[len(old):]
                    break
            normalized[new_key] = value
        return normalized

    def _load_checkpoint_in_py310(self, checkpoint_path: Path):
        fallback_python = Path(sys.executable)
        if not fallback_python.exists():
            raise HTTPException(
                status_code=500,
                detail=(
                    f"Checkpoint {checkpoint_path} needs fallback loader, but {fallback_python} does not exist."
                ),
            )

        with tempfile.NamedTemporaryFile(suffix=".pt", delete=False) as tmp:
            out_path = Path(tmp.name)

        script = (
            "import torch\n"
            "from pathlib import Path\n"
            f"src = Path(r'{checkpoint_path}')\n"
            f"dst = Path(r'{out_path}')\n"
            "obj = torch.load(src, map_location='cpu', weights_only=False)\n"
            "if isinstance(obj, dict) and 'state_dict' in obj:\n"
            "    obj = obj['state_dict']\n"
            "torch.save(obj, dst)\n"
            "print(dst)\n"
        )

        proc = subprocess.run(
            [str(fallback_python), "-c", script],
            capture_output=True,
            text=True,
            check=False,
        )
        if proc.returncode != 0:
            out_path.unlink(missing_ok=True)
            raise HTTPException(
                status_code=500,
                detail=(
                    f"Failed to convert checkpoint {checkpoint_path} with Python 3.10 fallback loader. stderr={proc.stderr.strip()}"
                ),
            )

        try:
            return torch.load(out_path, map_location="cpu", weights_only=False)
        finally:
            out_path.unlink(missing_ok=True)


model_manager = ModelManager()


class BBox(BaseModel):
    x1: float
    y1: float
    x2: float
    y2: float


class PersonSpec(BaseModel):
    bbox: BBox
    color: str = Field(default="#00ff00")


app = FastAPI(title="Gazelle+ Standalone Demo")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)
app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")


@app.on_event("startup")
def preload_default_model() -> None:
    if not PRELOAD_DEFAULT_MODEL:
        print(f"[startup] skipping default model preload: {DEFAULT_MODEL_NAME}")
        return
    print(f"[startup] preloading default model: {DEFAULT_MODEL_NAME}")
    model_manager.get(DEFAULT_MODEL_NAME)


def _clamp_bbox(bbox: BBox, width: int, height: int) -> tuple[int, int, int, int]:
    x1 = max(0, min(int(round(bbox.x1)), width - 1))
    y1 = max(0, min(int(round(bbox.y1)), height - 1))
    x2 = max(0, min(int(round(bbox.x2)), width))
    y2 = max(0, min(int(round(bbox.y2)), height))
    if x2 <= x1 or y2 <= y1:
        raise HTTPException(status_code=400, detail="Invalid bbox")
    return x1, y1, x2, y2


def _parse_hex_color(color: str) -> tuple[int, int, int, int]:
    text = (color or "#00ff00").strip()
    if text.startswith("#"):
        text = text[1:]
    if len(text) == 3:
        text = "".join(ch * 2 for ch in text)
    if len(text) != 6:
        return (0, 255, 0, 255)
    try:
        return (int(text[0:2], 16), int(text[2:4], 16), int(text[4:6], 16), 255)
    except ValueError:
        return (0, 255, 0, 255)


def _image_to_base64(image: Image.Image) -> str:
    import base64

    buffer = io.BytesIO()
    image.save(buffer, format="PNG")
    return "data:image/png;base64," + base64.b64encode(buffer.getvalue()).decode("utf-8")


def _downsample_for_inference(image: Image.Image) -> tuple[Image.Image, float]:
    max_side = max(image.size)
    if max_side <= MAX_INFER_SIDE:
        return image, 1.0
    scale = MAX_INFER_SIDE / max_side
    new_size = (max(1, int(round(image.width * scale))), max(1, int(round(image.height * scale))))
    return image.resize(new_size, Image.Resampling.LANCZOS), scale


def _query_gpu_status() -> dict:
    try:
        proc = subprocess.run(
            [
                "nvidia-smi",
                "--query-gpu=index,name,utilization.gpu,memory.used,memory.total",
                "--format=csv,noheader,nounits",
            ],
            capture_output=True,
            text=True,
            timeout=2,
            check=False,
        )
    except Exception as exc:
        return {"available": False, "message": str(exc), "gpus": []}

    if proc.returncode != 0:
        return {"available": False, "message": proc.stderr.strip() or "nvidia-smi unavailable", "gpus": []}

    gpus = []
    for line in proc.stdout.strip().splitlines():
        parts = [part.strip() for part in line.split(",")]
        if len(parts) < 5:
            continue
        index, name, util, mem_used, mem_total = parts[:5]
        try:
            gpus.append({
                "index": int(index),
                "name": name,
                "utilization": float(util),
                "memory_used": float(mem_used),
                "memory_total": float(mem_total),
            })
        except ValueError:
            continue
    return {"available": bool(gpus), "gpus": gpus}


def _make_head_inputs(image: Image.Image, bbox_px: tuple[int, int, int, int], head_transforms: list | None):
    if not head_transforms:
        return None
    crop = image.crop(bbox_px)
    return [transform(crop).unsqueeze(0) for transform in head_transforms]


def _find_peak(heatmap: torch.Tensor, width: int, height: int) -> tuple[int, int, float]:
    heatmap_np = heatmap.detach().cpu().numpy()
    flat_index = int(np.argmax(heatmap_np))
    peak_y, peak_x = np.unravel_index(flat_index, heatmap_np.shape)
    confidence = float(heatmap_np[peak_y, peak_x])
    target_x = int(round((peak_x / max(1, heatmap_np.shape[1] - 1)) * (width - 1)))
    target_y = int(round((peak_y / max(1, heatmap_np.shape[0] - 1)) * (height - 1)))
    return target_x, target_y, confidence


def _viz_sizes(image: Image.Image) -> dict[str, int]:
    short_side = max(1, min(image.size))
    scale = short_side / 512
    return {
        "bbox_width": max(2, int(round(4 * scale))),
        "line_width": max(2, int(round(5 * scale))),
        "point_radius": max(4, int(round(6 * scale))),
        "point_outline": max(1, int(round(1.5 * scale))),
    }


def _draw_result(image: Image.Image, bbox_px: tuple[int, int, int, int], target_xy: tuple[int, int], is_out: bool) -> Image.Image:
    canvas = image.copy().convert("RGBA")
    draw = ImageDraw.Draw(canvas)
    sizes = _viz_sizes(image)
    gaze_color = (0, 255, 0, 255)
    draw.rectangle(bbox_px, outline=gaze_color, width=sizes["bbox_width"])
    if not is_out:
        tx, ty = target_xy
        r = sizes["point_radius"]
        draw.ellipse((tx - r, ty - r, tx + r, ty + r), fill=gaze_color, outline=(255, 255, 255, 255), width=sizes["point_outline"])
        cx = (bbox_px[0] + bbox_px[2]) // 2
        cy = (bbox_px[1] + bbox_px[3]) // 2
        draw.line((cx, cy, tx, ty), fill=gaze_color, width=sizes["line_width"])
    else:
        draw.text((20, 20), "OUT", fill=(255, 64, 64, 255))
    return canvas.convert("RGB")


def _draw_out_of_frame_label(combined: Image.Image, bbox_px: tuple[int, int, int, int], color: tuple[int, int, int, int] = (255, 80, 80, 255)) -> None:
    draw = ImageDraw.Draw(combined)
    width, height = combined.size
    text = "OUT"
    font_path = "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf"
    max_text_w = max(40, int((bbox_px[2] - bbox_px[0]) * 1.5))
    max_text_h = max(18, int(height * 0.08))
    font_size = max(13, min(width, height) // 23)
    while font_size > 12:
        try:
            font = ImageFont.truetype(font_path, font_size)
        except Exception:
            font = ImageFont.load_default()
            break
        bbox = draw.textbbox((0, 0), text, font=font)
        text_w = bbox[2] - bbox[0]
        text_h = bbox[3] - bbox[1]
        if text_w <= max_text_w and text_h <= max_text_h:
            break
        font_size -= 2
    bbox = draw.textbbox((0, 0), text, font=font)
    text_w = bbox[2] - bbox[0]
    text_h = bbox[3] - bbox[1]
    text_offset_x = bbox[0]
    text_offset_y = bbox[1]
    pad_x = max(10, font_size // 3)
    pad_y = max(12, font_size // 2)
    box_w = text_w + 2 * pad_x
    box_h = text_h + 2 * pad_y
    x0 = max(0, min(width - box_w, (bbox_px[0] + bbox_px[2] - box_w) // 2))
    y0 = min(height - box_h, bbox_px[3] + max(3, pad_y // 3))
    if y0 < 0 or y0 <= bbox_px[3] - 1:
        y0 = max(0, bbox_px[1] - box_h - max(3, pad_y // 3))
    draw.rounded_rectangle(
        (x0, y0, x0 + box_w, y0 + box_h),
        radius=max(4, min(pad_x, pad_y)),
        fill=(0, 0, 0, 140),
        outline=color[:3] + (230,),
        width=max(2, min(pad_x, pad_y) // 3),
    )
    text_x = x0 + (box_w - text_w) / 2 - text_offset_x
    text_y = y0 + (box_h - text_h) / 2 - text_offset_y - max(0.5, font_size * 0.04)
    draw.text((text_x, text_y), text, fill=color, font=font)


def _draw_combined_visualization(
    image: Image.Image,
    bbox_px: tuple[int, int, int, int],
    bbox_norm: list[float],
    heatmap: torch.Tensor,
    target_xy: tuple[int, int],
    is_out: bool,
    show_heatmap: bool = True,
    show_gaze: bool = True,
) -> Image.Image:
    render_scale = max(1, int(np.ceil(768 / max(1, min(image.size)))))
    if render_scale > 1:
        image = image.resize((image.width * render_scale, image.height * render_scale), Image.Resampling.LANCZOS)
        bbox_px = tuple(int(round(v * render_scale)) for v in bbox_px)
        bbox_norm = [bbox_px[0] / image.width, bbox_px[1] / image.height, bbox_px[2] / image.width, bbox_px[3] / image.height]
        target_xy = (int(round(target_xy[0] * render_scale)), int(round(target_xy[1] * render_scale)))

    if is_out or not show_heatmap:
        combined = image.copy().convert("RGBA")
    else:
        combined = visualize_heatmap(image, heatmap, bbox=bbox_norm).convert("RGBA")
    draw = ImageDraw.Draw(combined)
    sizes = _viz_sizes(image)
    gaze_color = (0, 255, 0, 255)
    draw.rectangle(bbox_px, outline=gaze_color, width=sizes["bbox_width"])
    if is_out:
        _draw_out_of_frame_label(combined, bbox_px)
    elif not show_gaze:
        pass
    else:
        tx, ty = target_xy
        r = sizes["point_radius"]
        draw.ellipse((tx - r, ty - r, tx + r, ty + r), fill=gaze_color, outline=(255, 255, 255, 255), width=sizes["point_outline"])
        cx = (bbox_px[0] + bbox_px[2]) // 2
        cy = (bbox_px[1] + bbox_px[3]) // 2
        draw.line((cx, cy, tx, ty), fill=gaze_color, width=sizes["line_width"])
    return combined.convert("RGB")


def _draw_multi_visualization(
    image: Image.Image,
    bbox_px_list: list[tuple[int, int, int, int]],
    bbox_norm_list: list[list[float]],
    heatmaps: list[torch.Tensor],
    target_xy_list: list[tuple[int, int]],
    is_out_list: list[bool],
    colors: list[str],
    show_heatmap: bool = False,
    show_gaze: bool = True,
) -> Image.Image:
    render_scale = max(1, int(np.ceil(768 / max(1, min(image.size)))))
    if render_scale > 1:
        image = image.resize((image.width * render_scale, image.height * render_scale), Image.Resampling.LANCZOS)
        bbox_px_list = [tuple(int(round(v * render_scale)) for v in bbox) for bbox in bbox_px_list]
        bbox_norm_list = [[bbox[0] / image.width, bbox[1] / image.height, bbox[2] / image.width, bbox[3] / image.height] for bbox in bbox_px_list]
        target_xy_list = [(int(round(x * render_scale)), int(round(y * render_scale))) for x, y in target_xy_list]

    combined = image.copy().convert("RGBA")
    for heatmap, bbox_norm, is_out in zip(heatmaps, bbox_norm_list, is_out_list):
        if not show_heatmap:
            continue
        if is_out:
            continue
        overlay = visualize_heatmap(image, heatmap, bbox=None).convert("RGBA")
        combined = Image.blend(combined, overlay, alpha=0.45)

    draw = ImageDraw.Draw(combined)
    sizes = _viz_sizes(image)
    for bbox_px, target_xy, is_out, color in zip(bbox_px_list, target_xy_list, is_out_list, colors):
        rgba = _parse_hex_color(color)
        draw.rectangle(bbox_px, outline=rgba, width=sizes["bbox_width"])
        if is_out:
            _draw_out_of_frame_label(combined, bbox_px, rgba)
            continue
        if not show_gaze:
            continue
        tx, ty = target_xy
        r = sizes["point_radius"]
        draw.ellipse((tx - r, ty - r, tx + r, ty + r), fill=rgba, outline=(255, 255, 255, 255), width=sizes["point_outline"])
        cx = (bbox_px[0] + bbox_px[2]) // 2
        cy = (bbox_px[1] + bbox_px[3]) // 2
        draw.line((cx, cy, tx, ty), fill=rgba, width=sizes["line_width"])
    return combined.convert("RGB")


@app.get("/")
def index():
    return FileResponse(str(STATIC_DIR / "index.html"))


@app.get("/api/models")
def list_models():
    def serialize_model(name: str, cfg: dict) -> dict:
        item = {
            "name": name,
            "source": cfg.get("source", "local_checkpoint"),
            "display_name": cfg.get("display_name", name),
        }
        if cfg.get("source") == "hf_page":
            item.update({
                "repo_id": cfg["repo_id"],
                "requires_token": True,
                "token_env_vars": list(HF_TOKEN_ENV_VARS),
            })
        else:
            item.update({
                "checkpoint": str(cfg["checkpoint"]),
                "factory_name": cfg["factory_name"],
                "include_backbone": cfg["include_backbone"],
                "resolved_backbone_paths": _resolve_backbone_paths(cfg["factory_name"]),
                "env_hints": MODEL_PATH_ENV[cfg["factory_name"]],
            })
        return item

    return {
        "models": [serialize_model(name, cfg) for name, cfg in MODEL_OPTIONS.items()]
    }


@app.get("/api/gpu")
def gpu_status():
    return _query_gpu_status()


@app.post("/api/preload")
async def preload_model(model_name: str = Form(...)):
    model_manager.get(model_name)
    return {"ok": True, "model_name": model_name}


@app.post("/api/infer")
async def infer(
    model_name: str = Form(...),
    bbox: str = Form(...),
    image: UploadFile = File(...),
    show_heatmap: bool = Form(True),
    show_gaze: bool = Form(True),
):
    try:
        bbox_obj = BBox.model_validate(json.loads(bbox))
    except Exception as exc:
        raise HTTPException(status_code=400, detail=f"Invalid request payload: {exc}") from exc

    raw = await image.read()
    if not raw:
        raise HTTPException(status_code=400, detail="Empty image")

    pil_image = Image.open(io.BytesIO(raw)).convert("RGB")
    original_width, original_height = pil_image.size
    pil_image, image_scale = _downsample_for_inference(pil_image)
    width, height = pil_image.size
    if image_scale != 1.0:
        bbox_obj = BBox(x1=bbox_obj.x1 * image_scale, y1=bbox_obj.y1 * image_scale, x2=bbox_obj.x2 * image_scale, y2=bbox_obj.y2 * image_scale)
    bbox_px = _clamp_bbox(bbox_obj, width, height)
    bbox_norm = [bbox_px[0] / width, bbox_px[1] / height, bbox_px[2] / width, bbox_px[3] / height]

    model, scene_transforms, head_transforms = model_manager.get(model_name)
    scene_inputs = [transform(pil_image).unsqueeze(0).to(model_manager.device) for transform in scene_transforms]
    head_inputs = _make_head_inputs(pil_image, bbox_px, head_transforms)
    if head_inputs is not None:
        head_inputs = [tensor.to(model_manager.device) for tensor in head_inputs]

    with torch.inference_mode():
        output = model({
            "images": scene_inputs,
            "head_images": head_inputs,
            "bboxes": [[tuple(bbox_norm)]],
        })

    heatmap = output["heatmap"][0][0]
    inout_score = float(output["inout"][0][0].detach().cpu().item()) if output.get("inout") is not None else 1.0
    is_out = inout_score < OUT_THRESHOLD
    target_x, target_y, peak_score = _find_peak(heatmap, width, height)

    heatmap_overlay = visualize_heatmap(pil_image, heatmap, bbox=bbox_norm).convert("RGB")
    result_overlay = _draw_result(pil_image, bbox_px, (target_x, target_y), is_out)
    combined_overlay = _draw_combined_visualization(pil_image, bbox_px, bbox_norm, heatmap, (target_x, target_y), is_out, show_heatmap=show_heatmap, show_gaze=show_gaze)

    return {
        "model_name": model_name,
        "bbox": {"x1": bbox_px[0], "y1": bbox_px[1], "x2": bbox_px[2], "y2": bbox_px[3]},
        "target": None if is_out else {"x": target_x, "y": target_y, "score": peak_score},
        "inout_score": inout_score,
        "is_out": is_out,
        "heatmap_image": _image_to_base64(heatmap_overlay),
        "result_image": _image_to_base64(result_overlay),
        "combined_image": _image_to_base64(combined_overlay),
        "image_size": {"width": width, "height": height},
        "original_image_size": {"width": original_width, "height": original_height},
    }


@app.post("/api/infer_multi")
async def infer_multi(
    model_name: str = Form(...),
    people: str = Form(...),
    image: UploadFile = File(...),
    show_heatmap: bool = Form(False),
    show_gaze: bool = Form(True),
):
    try:
        raw_people = json.loads(people)
        person_specs = [PersonSpec.model_validate(item) for item in raw_people]
    except Exception as exc:
        raise HTTPException(status_code=400, detail=f"Invalid people payload: {exc}") from exc
    if not person_specs:
        raise HTTPException(status_code=400, detail="No bboxes provided")

    raw = await image.read()
    if not raw:
        raise HTTPException(status_code=400, detail="Empty image")

    pil_image = Image.open(io.BytesIO(raw)).convert("RGB")
    original_width, original_height = pil_image.size
    pil_image, image_scale = _downsample_for_inference(pil_image)
    width, height = pil_image.size
    model, scene_transforms, head_transforms = model_manager.get(model_name)
    scene_inputs = [transform(pil_image).unsqueeze(0).to(model_manager.device) for transform in scene_transforms]

    results = []
    bbox_px_list = []
    bbox_norm_list = []
    heatmaps = []
    target_xy_list = []
    is_out_list = []
    colors = []

    for idx, spec in enumerate(person_specs):
        bbox_obj = spec.bbox
        if image_scale != 1.0:
            bbox_obj = BBox(x1=bbox_obj.x1 * image_scale, y1=bbox_obj.y1 * image_scale, x2=bbox_obj.x2 * image_scale, y2=bbox_obj.y2 * image_scale)
        bbox_px = _clamp_bbox(bbox_obj, width, height)
        bbox_norm = [bbox_px[0] / width, bbox_px[1] / height, bbox_px[2] / width, bbox_px[3] / height]
        head_inputs = _make_head_inputs(pil_image, bbox_px, head_transforms)
        if head_inputs is not None:
            head_inputs = [tensor.to(model_manager.device) for tensor in head_inputs]
        with torch.inference_mode():
            output = model({
                "images": scene_inputs,
                "head_images": head_inputs,
                "bboxes": [[tuple(bbox_norm)]],
            })
        heatmap = output["heatmap"][0][0]
        inout_score = float(output["inout"][0][0].detach().cpu().item()) if output.get("inout") is not None else 1.0
        is_out = inout_score < OUT_THRESHOLD
        target_x, target_y, peak_score = _find_peak(heatmap, width, height)

        bbox_px_list.append(bbox_px)
        bbox_norm_list.append(bbox_norm)
        heatmaps.append(heatmap)
        target_xy_list.append((target_x, target_y))
        is_out_list.append(is_out)
        colors.append(spec.color)
        results.append({
            "index": idx,
            "color": spec.color,
            "bbox": {"x1": bbox_px[0], "y1": bbox_px[1], "x2": bbox_px[2], "y2": bbox_px[3]},
            "target": None if is_out else {"x": target_x, "y": target_y, "score": peak_score},
            "inout_score": inout_score,
            "is_out": is_out,
        })

    combined_overlay = _draw_multi_visualization(pil_image, bbox_px_list, bbox_norm_list, heatmaps, target_xy_list, is_out_list, colors, show_heatmap=show_heatmap, show_gaze=show_gaze)
    return {
        "model_name": model_name,
        "people": results,
        "combined_image": _image_to_base64(combined_overlay),
        "image_size": {"width": width, "height": height},
        "original_image_size": {"width": original_width, "height": original_height},
    }


if __name__ == "__main__":
    import uvicorn

    uvicorn.run("app:app", host="0.0.0.0", port=int(os.environ.get("PORT", "7860")), reload=False)