"""HuggingFace Space — SAM Assist for UCL AI VIM.

CPU-Basic Space (2 vCPU, 16 GB RAM, no GPU). All inference runs on CPU; SAM 2
tiny + MedSAM 2 + Grounding DINO tiny are the only variants light enough to be
usable at this size. Video tracking is intentionally not exposed here because
SAM2VideoPredictor + a 60–120 s HF call timeout combine into something that
never finishes.

Endpoints (POST to https://<owner>-<space>.hf.space/run/<api_name>):
    api_name="segment_point"   -> (model, frame_b64, x, y, label, class_id, variant) -> polys json
    api_name="segment_text"    -> (frame_b64, text, class_id, box_thr, text_thr, variant) -> polys json
    api_name="classes"         -> () -> taxonomy json
    api_name="healthz"         -> () -> {"ok": True, ...}

The React SamBridge picks the right call shape based on VITE_SAM_API_TYPE.
"""
from __future__ import annotations
import base64
import io
import json
import logging
import os
from typing import Optional

import cv2
import gradio as gr
import numpy as np
import torch
from PIL import Image

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s %(message)s")
log = logging.getLogger("sam-hf")

# Force CPU — this Space has no GPU and we don't want a runtime error from
# anything reaching for cuda.
DEVICE = "cpu"
torch.set_num_threads(max(1, (os.cpu_count() or 2) - 1))

# --- taxonomy --------------------------------------------------------------
TAXONOMY = {
    "regions": [
        {"id": "grasper",           "name": "Grasper",           "color": "#4499cc"},
        {"id": "scissors",          "name": "Scissors",          "color": "#cc4444"},
        {"id": "clip_applier",      "name": "Clip Applier",      "color": "#ccaa33"},
        {"id": "monopolar_cautery", "name": "Monopolar Cautery", "color": "#cc7733"},
        {"id": "suction",           "name": "Suction/Irrigator", "color": "#33997a"},
        {"id": "needle_driver",     "name": "Needle Driver",     "color": "#8844cc"},
        {"id": "trocar",            "name": "Trocar",            "color": "#5577cc"},
        {"id": "kidney",            "name": "Kidney",            "color": "#dd8855"},
        {"id": "tumour",            "name": "Tumour",            "color": "#ff5566"},
        {"id": "renal_artery",      "name": "Renal Artery",      "color": "#ee3344"},
        {"id": "renal_vein",        "name": "Renal Vein",        "color": "#4466dd"},
    ],
    "events": [
        {"id": "bleeding",        "name": "Bleeding",            "color": "#ff2233"},
        {"id": "smoke",           "name": "Smoke/Fog",           "color": "#999999"},
        {"id": "clipping",        "name": "Clipping",            "color": "#ccaa33"},
        {"id": "cutting",         "name": "Cutting",             "color": "#cc4444"},
        {"id": "clamp_on",        "name": "Clamp-on (WIT start)","color": "#33ccff"},
        {"id": "clamp_off",       "name": "Clamp-off (WIT end)", "color": "#33cc66"},
        {"id": "needle_delivery", "name": "Needle delivery",     "color": "#cc99ff"},
        {"id": "suture_delivery", "name": "Suture delivery",     "color": "#9966ff"},
        {"id": "suction_response","name": "Suction response",    "color": "#33aaaa"},
        {"id": "occlusion",       "name": "Assistant occlusion", "color": "#999933"},
        {"id": "wrong_item",      "name": "Wrong-item / reload", "color": "#ff7733"},
        {"id": "idle_with_need",  "name": "Idle-with-need (>=2s)","color": "#bbbbbb"},
    ],
    "phases": [
        {"id": "phase_setup",         "name": "Setup / port placement",                     "color": "#5b8def"},
        {"id": "phase_colon_mob",     "name": "Colon mobilisation & retroperitoneal access","color": "#6fa8dc"},
        {"id": "phase_hilar_dissect", "name": "Hilar dissection",                           "color": "#7d6cb1"},
        {"id": "phase_tumour_id",     "name": "Tumour identification",                      "color": "#a364b1"},
        {"id": "phase_hilar_clamp",   "name": "Hilar clamping",                             "color": "#d65a8c"},
        {"id": "phase_resection",     "name": "Tumour resection / enucleation",             "color": "#e94f64"},
        {"id": "phase_renorrhaphy",   "name": "Renorrhaphy",                                "color": "#e87f3f"},
        {"id": "phase_unclamp",       "name": "Unclamping & haemostasis",                   "color": "#e0a93b"},
        {"id": "phase_specimen",      "name": "Specimen retrieval & closure",               "color": "#7fb069"},
        {"id": "phase_idle",          "name": "Out-of-body / idle",                         "color": "#888888"},
    ],
}

# Only the smallest SAM 2 variant is realistically usable on CPU.
SAM2_VARIANT_IDS = {"tiny": "facebook/sam2.1-hiera-tiny"}


# --- shared utilities ------------------------------------------------------

def _decode_b64(b64: str) -> np.ndarray:
    if "," in b64:
        b64 = b64.split(",", 1)[1]
    return np.array(Image.open(io.BytesIO(base64.b64decode(b64))).convert("RGB"))


def _mask_to_polys(mask: np.ndarray, eps: float = 1.0) -> list[list[list[float]]]:
    m = (mask.astype(bool).astype(np.uint8)) * 255
    contours, _ = cv2.findContours(m, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    out: list[list[list[float]]] = []
    for c in contours:
        if cv2.contourArea(c) < 10:
            continue
        approx = cv2.approxPolyDP(c, eps, True)
        if len(approx) < 3:
            continue
        out.append([[float(p[0][0]), float(p[0][1])] for p in approx])
    return out


# --- lazy model holders ----------------------------------------------------
# Loaded on first call, then cached for the lifetime of the process.
_SAM2 = {"p": None}
_MEDSAM2 = {"p": None}
_GD = {"processor": None, "model": None}


def _load_sam2():
    from sam2.sam2_image_predictor import SAM2ImagePredictor
    if _SAM2["p"] is None:
        log.info("Loading SAM 2 tiny (CPU)...")
        _SAM2["p"] = SAM2ImagePredictor.from_pretrained(SAM2_VARIANT_IDS["tiny"], device=DEVICE)
    return _SAM2["p"]


def _load_medsam2():
    from sam2.sam2_image_predictor import SAM2ImagePredictor
    if _MEDSAM2["p"] is None:
        log.info("Loading MedSAM 2 (wanglab/MedSAM2, CPU)...")
        _MEDSAM2["p"] = SAM2ImagePredictor.from_pretrained("wanglab/MedSAM2", device=DEVICE)
    return _MEDSAM2["p"]


def _load_gd():
    from transformers import AutoProcessor, AutoModelForZeroShotObjectDetection
    if _GD["model"] is None:
        log.info("Loading Grounding DINO tiny (CPU)...")
        _GD["processor"] = AutoProcessor.from_pretrained("IDEA-Research/grounding-dino-tiny")
        _GD["model"] = (
            AutoModelForZeroShotObjectDetection.from_pretrained("IDEA-Research/grounding-dino-tiny")
            .to(DEVICE).eval()
        )
    return _GD["processor"], _GD["model"]


# --- endpoints --------------------------------------------------------------

def healthz() -> str:
    return json.dumps({
        "ok": True,
        "runtime": "hf-cpu-basic",
        "device": DEVICE,
        "torch": torch.__version__,
        "cpu_threads": torch.get_num_threads(),
    })


def classes() -> str:
    return json.dumps(TAXONOMY)


def segment_point(model: str, frame_b64: str, x: float, y: float,
                  label: int = 1, class_id: Optional[str] = None,
                  variant: str = "tiny") -> str:
    img = _decode_b64(frame_b64)
    h, w = img.shape[:2]

    if model == "sam2":
        predictor = _load_sam2()
    elif model == "medsam2":
        predictor = _load_medsam2()
    elif model == "grounded_sam2":
        # text-only model; fall back to SAM 2 for point prompts so the UI can
        # keep "grounded_sam2" selected while toggling prompt type
        predictor = _load_sam2()
    else:
        return json.dumps({"error": f"unknown model {model!r}", "polygons": [], "width": w, "height": h})

    predictor.set_image(img)
    pts = np.array([[x, y]], dtype=np.float32)
    lbs = np.array([int(label)], dtype=np.int32)
    with torch.inference_mode():
        masks, scores, _ = predictor.predict(point_coords=pts, point_labels=lbs, multimask_output=True)
    best = int(np.argmax(scores))
    polys = [
        {"points": pts_xy, "score": float(scores[best]), "label": class_id}
        for pts_xy in _mask_to_polys(masks[best])
    ]
    return json.dumps({"polygons": polys, "width": w, "height": h})


def segment_text(frame_b64: str, text: str, class_id: Optional[str] = None,
                 box_threshold: float = 0.30, text_threshold: float = 0.25,
                 variant: str = "tiny") -> str:
    img = _decode_b64(frame_b64)
    h, w = img.shape[:2]

    processor, gd_model = _load_gd()
    sam2 = _load_sam2()

    prompt = text.strip().lower()
    if not prompt.endswith("."):
        prompt += "."

    pil = Image.fromarray(img)
    inputs = processor(images=pil, text=prompt, return_tensors="pt").to(DEVICE)
    with torch.inference_mode():
        outputs = gd_model(**inputs)
    results = processor.post_process_grounded_object_detection(
        outputs, inputs.input_ids,
        box_threshold=box_threshold, text_threshold=text_threshold,
        target_sizes=[(h, w)],
    )[0]
    boxes = results["boxes"].detach().cpu().numpy()
    scores = results["scores"].detach().cpu().numpy()
    phrases = results["labels"]

    polys: list[dict] = []
    if len(boxes):
        sam2.set_image(img)
        with torch.inference_mode():
            for box, score, phrase in zip(boxes, scores, phrases):
                masks, _, _ = sam2.predict(box=box, multimask_output=False)
                for ring in _mask_to_polys(masks[0]):
                    polys.append({"points": ring, "score": float(score), "label": str(phrase or class_id)})
    return json.dumps({"polygons": polys, "width": w, "height": h})


# --- Gradio Blocks ----------------------------------------------------------

with gr.Blocks(title="UCL AI VIM · SAM Assist (CPU)") as demo:
    gr.Markdown(
        "## UCL AI VIM · SAM Assist (HF Space, CPU Basic)\n"
        "Backend for the [annotation app](https://syncsurge.netlify.app/). Three models, point + text prompts.\n"
        "\n"
        "> ⚠️ **CPU-only Space — point prompts take ~10–30 s, text prompts ~30–60 s.** First call after "
        "a cold start is slower (weights download). For interactive speed, run the FastAPI backend "
        "locally against your own GPU."
    )
    with gr.Tab("segment-point"):
        in_model    = gr.Dropdown(choices=["sam2", "medsam2", "grounded_sam2"], value="sam2", label="model")
        in_frame    = gr.Textbox(label="frame_b64 (data:image/jpeg;base64,...)", lines=2)
        in_x        = gr.Number(label="x", value=0.0)
        in_y        = gr.Number(label="y", value=0.0)
        in_label    = gr.Number(label="label (1=fg, 0=bg)", value=1, precision=0)
        in_class    = gr.Textbox(label="class_id (optional)")
        in_variant  = gr.Dropdown(choices=list(SAM2_VARIANT_IDS), value="tiny", label="variant")
        out_point   = gr.Textbox(label="polygons json")
        gr.Button("run").click(
            fn=segment_point,
            inputs=[in_model, in_frame, in_x, in_y, in_label, in_class, in_variant],
            outputs=out_point, api_name="segment_point",
        )
    with gr.Tab("segment-text"):
        tx_frame    = gr.Textbox(label="frame_b64", lines=2)
        tx_text     = gr.Textbox(label="text prompt", value="kidney tumour")
        tx_class    = gr.Textbox(label="class_id (optional)")
        tx_box_thr  = gr.Slider(0.05, 0.9, value=0.30, label="box_threshold")
        tx_text_thr = gr.Slider(0.05, 0.9, value=0.25, label="text_threshold")
        tx_variant  = gr.Dropdown(choices=list(SAM2_VARIANT_IDS), value="tiny", label="variant")
        tx_out      = gr.Textbox(label="polygons json")
        gr.Button("run").click(
            fn=segment_text,
            inputs=[tx_frame, tx_text, tx_class, tx_box_thr, tx_text_thr, tx_variant],
            outputs=tx_out, api_name="segment_text",
        )
    with gr.Tab("meta"):
        gr.Button("/healthz").click(fn=healthz, outputs=gr.Textbox(label="health"), api_name="healthz")
        gr.Button("/classes").click(fn=classes, outputs=gr.Textbox(label="classes"), api_name="classes")


if __name__ == "__main__":
    # HF Spaces auto-detects host/port from environment. Passing them
    # explicitly can cause "localhost not accessible" errors in some
    # Space configurations.
    demo.queue(max_size=10).launch()