"""SAM 3 video concept-tracking API (ZeroGPU), transformers Sam3VideoModel route.

Tracks every instance of the given concept(s) across video frames with stable
object ids. Output schema matches the local video_client parser:
  {version, model, fps, width, height, n_frames, tracks:[{label, object_id,
   frames:[{frame, score, box, mask_png_b64}]}]}
"""
import base64
import io
import os

import gradio as gr
import numpy as np
import spaces
from PIL import Image
from transformers import Sam3VideoModel, Sam3VideoProcessor

HF_TOKEN = os.environ.get("HF_TOKEN")
MODEL_ID = "facebook/sam3"

# Built at import on CPU; moved to CUDA inside the @spaces.GPU function.
processor = Sam3VideoProcessor.from_pretrained(MODEL_ID, token=HF_TOKEN)
model = Sam3VideoModel.from_pretrained(MODEL_ID, token=HF_TOKEN)
model.eval()


def _enc(mask_bool: np.ndarray, maxside: int = 512) -> str:
    h, w = mask_bool.shape
    img = Image.fromarray((mask_bool.astype(np.uint8) * 255), "L")
    scale = min(1.0, maxside / max(h, w))
    if scale < 1.0:
        img = img.resize((max(1, int(w * scale)), max(1, int(h * scale))))
    buf = io.BytesIO(); img.save(buf, "PNG")
    return base64.b64encode(buf.getvalue()).decode("ascii")


def _np(x):
    return x.detach().cpu().numpy() if hasattr(x, "detach") else np.asarray(x)


def _read_frames(path, max_frames):
    """Sample up to `max_frames` frames EVENLY across the whole video (proper intervals), not just the
    first N. A clip with <= max_frames frames is taken in full; a longer clip is sub-sampled at a
    constant stride so the selection spans start→end."""
    import imageio
    max_frames = max(1, int(max_frames))
    reader = imageio.get_reader(path)
    try:
        try:
            n = int(reader.count_frames())
        except Exception:  # some streams can't report a count -> fall back to sequential
            n = 0
        if n > max_frames:
            # evenly-spaced indices spanning 0 .. n-1 (e.g. 96 frames, max 24 -> every 4th frame)
            idxs = sorted({round(i * (n - 1) / (max_frames - 1)) for i in range(max_frames)}) \
                if max_frames > 1 else [0]
            frames = []
            for i in idxs:
                try:
                    frames.append(np.asarray(reader.get_data(i)))
                except Exception:
                    break
            if frames:
                return frames
        # short clip (or no count / seek unsupported): read sequentially up to max_frames
        frames = []
        for i, fr in enumerate(reader):
            if i >= max_frames:
                break
            frames.append(np.asarray(fr))
        return frames
    finally:
        reader.close()


@spaces.GPU(duration=300)
def api_track(video, concepts, conf, max_frames):
    """Streaming generator: yields {done:False, progress, desc} per frame, then a
    final {done:True, ..., tracks:[...]}. (gr.Progress can't cross ZeroGPU's process
    boundary, so we stream progress as output instead.)"""
    device = "cuda"
    model.to(device)
    concept_list = [c.strip() for c in (concepts or "").split(",") if c.strip()] or ["person"]
    frames = _read_frames(video, int(max_frames))
    if not frames:
        yield {"done": True, "error": "no frames read from video", "tracks": []}
        return
    H, W = frames[0].shape[:2]
    total = max(1, min(len(frames), int(max_frames)))
    yield {"done": False, "progress": 0.0, "desc": f"loaded {len(frames)} frames; starting tracker"}
    session = processor.init_video_session(
        video=frames, inference_device=device,
        processing_device="cpu", video_storage_device="cpu",
    )
    processor.add_text_prompt(session, concept_list)

    tracks, obj_label, n_frames = {}, {}, 0
    for mo in model.propagate_in_video_iterator(inference_session=session,
                                                max_frame_num_to_track=int(max_frames)):
        proc = processor.postprocess_outputs(session, mo)
        fi = int(mo.frame_idx); n_frames = max(n_frames, fi + 1)
        for prompt, oids in (proc.get("prompt_to_obj_ids") or {}).items():
            for oid in oids:
                obj_label.setdefault(int(oid), prompt)
        oids = _np(proc["object_ids"]).tolist()
        scores = _np(proc["scores"]).tolist()
        masks = proc["masks"]
        boxes = _np(proc["boxes"])
        for k, oid in enumerate(oids):
            oid = int(oid)
            m = _np(masks[k])
            if m.ndim == 3:
                m = m[0]
            m = m > 0.5 if m.dtype != bool else m
            tr = tracks.get(oid)
            if tr is None:
                tr = {"label": obj_label.get(oid, concept_list[0]), "object_id": oid, "frames": []}
                tracks[oid] = tr
            tr["frames"].append({"frame": fi, "score": float(scores[k]),
                                 "box": [float(v) for v in boxes[k]],
                                 "mask_png_b64": _enc(m)})
        yield {"done": False, "progress": min(fi + 1, total) / total,
               "desc": f"frame {fi + 1}/{total}"}

    out_tracks = []
    for oid, tr in tracks.items():
        tr["label"] = obj_label.get(oid, tr["label"])
        if tr["frames"] and max(f["score"] for f in tr["frames"]) >= float(conf):
            out_tracks.append(tr)
    yield {"done": True, "version": "3", "model": MODEL_ID, "fps": 0.0,
           "width": W, "height": H, "n_frames": n_frames, "tracks": out_tracks}


with gr.Blocks(title="SAM3 Video") as demo:
    gr.Markdown("# SAM 3 Video Tracking API\nUpload a video, enter comma-separated concepts.")
    with gr.Row():
        vid = gr.File(file_count="single", type="filepath", label="Video (mp4)")
        out = gr.JSON(label="Tracks")
    txt = gr.Textbox(label="Concepts (comma-separated)", value="person")
    conf = gr.Slider(0.0, 1.0, value=0.4, step=0.05, label="Confidence")
    mf = gr.Slider(8, 96, value=48, step=8, label="Max frames")
    gr.Button("Track").click(api_track, [vid, txt, conf, mf], out, api_name="api_track")

if __name__ == "__main__":
    demo.queue().launch(show_error=True)