import os
import shutil
import tempfile

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # MUST come before torch / any CUDA-touching import
import torch
import torch.nn.functional as F
import numpy as np
import cv2
from PIL import Image
import gradio as gr
from huggingface_hub import hf_hub_download
from gradio_client import Client, handle_file
from iopath.common.file_io import g_pathmgr

# Patch torch.load to allow weights_only=False for checkpoints
_orig_load = torch.load


def _patched_load(*args, **kwargs):
    kwargs["weights_only"] = kwargs.get("weights_only", False)
    return _orig_load(*args, **kwargs)


torch.load = _patched_load

from sam2.build_sam import build_sam2matting, build_sam2matting_video_predictor
from sam2.sam2matting_image_predictor import SAM2MattingImagePredictor

MODEL_ID = "FudanCVL/SAM2Matting"
REMBG_SPACE_ID = "gokaygokay/Inspyrenet-Rembg"
DEVICE = "cuda"

COLOR_MAP = {
    "Green": [0, 255, 0],
    "Blue": [0, 0, 255],
    "Red": [255, 0, 0],
    "White": [255, 255, 255],
    "Black": [0, 0, 0],
    "Gray": [128, 128, 128],
}

# variant -> (family, ckpt filename, config or None, mask_prompt_size)
VARIANTS = {
    "SAM2.1-Tiny": (
        "sam2",
        "checkpoints/SAM2Matting-SAM2.1Tiny.pt",
        "configs/sam2matting-sam2.1tiny.yaml",
        256,
    ),
    "SAM2.1-Base+": (
        "sam2",
        "checkpoints/SAM2Matting-SAM2.1Base+.pt",
        "configs/sam2matting-sam2.1base+.yaml",
        256,
    ),
    "SAM3": (
        "sam3",
        "checkpoints/SAM2Matting-SAM3.pt",
        None,
        288,
    ),
}

_ckpt_cache: dict[str, str] = {}
_image_predictors: dict[str, object] = {}
_video_predictors: dict[str, object] = {}


def _download_ckpt(filename: str) -> str:
    if filename not in _ckpt_cache:
        _ckpt_cache[filename] = hf_hub_download(
            repo_id=MODEL_ID, filename=filename, repo_type="model"
        )
    return _ckpt_cache[filename]


def _load_sam3_tracker_state_dict(checkpoint: str) -> dict:
    with g_pathmgr.open(checkpoint, "rb") as f:
        ckpt = torch.load(f, map_location="cpu", weights_only=True)
    sd = ckpt["model"]
    out = {}
    for k, v in sd.items():
        if k.startswith("detector.backbone.vision_backbone."):
            out[k.removeprefix("detector.")] = v
        elif k.startswith("tracker."):
            out[k.removeprefix("tracker.")] = v
    return out


def get_image_predictor(variant: str):
    if variant in _image_predictors:
        return _image_predictors[variant], VARIANTS[variant][3]

    family, ckpt_name, config, mask_size = VARIANTS[variant]
    ckpt_path = _download_ckpt(ckpt_name)

    if family == "sam2":
        model = build_sam2matting(
            config_file=config, ckpt_path=ckpt_path, device=DEVICE
        )
        predictor = SAM2MattingImagePredictor(model)
    else:
        from sam3.model.build_sam3matting import build_sam3matting
        from sam3.model.sam3matting_image_predictor import SAM3MattingImagePredictor

        sd = _load_sam3_tracker_state_dict(ckpt_path)
        model = build_sam3matting(checkpoint=None, device=DEVICE)
        model.load_state_dict(sd, strict=False)
        predictor = SAM3MattingImagePredictor(model)

    _image_predictors[variant] = predictor
    return predictor, mask_size


def get_video_predictor(variant: str):
    if variant in _video_predictors:
        return _video_predictors[variant], VARIANTS[variant][3]

    family, ckpt_name, config, mask_size = VARIANTS[variant]
    ckpt_path = _download_ckpt(ckpt_name)

    if family == "sam2":
        predictor = build_sam2matting_video_predictor(
            config, ckpt_path, device=DEVICE
        )
    else:
        from sam3.model.sam3matting_video_predictor import (
            build_sam3matting_video_predictor,
        )

        sd = _load_sam3_tracker_state_dict(ckpt_path)
        predictor = build_sam3matting_video_predictor(
            checkpoint=None, device=DEVICE
        )
        predictor.load_state_dict(sd, strict=False)

    _video_predictors[variant] = predictor
    return predictor, mask_size


def _auto_generate_mask(image: Image.Image) -> Image.Image:
    tmp_path = os.path.join("/tmp", "rembg_input.png")
    image.save(tmp_path)
    client = Client(REMBG_SPACE_ID, token=os.environ.get("HF_TOKEN"))
    mask_path = client.predict(
        input_image=handle_file(tmp_path),
        output_type="Mask only",
        api_name="/predict",
    )
    return Image.open(mask_path).convert("L")


def _prepare_image_mask_tensors(mask: Image.Image, mask_size: int):
    mask_np = np.array(mask.convert("L"))
    raw_mask = (torch.from_numpy(mask_np) / 255) > 0
    mask_input = (torch.from_numpy(mask_np) > 0).float() * 20 - 10
    mask_input = mask_input.unsqueeze(0).unsqueeze(0)
    mask_input = F.interpolate(
        mask_input,
        size=(mask_size, mask_size),
        mode="bilinear",
        align_corners=False,
    )
    return raw_mask, mask_input


def _compose_rgb(orig_rgb: np.ndarray, alpha_u8: np.ndarray, bg_color: str) -> np.ndarray:
    bg = np.full(
        (*alpha_u8.shape, 3), COLOR_MAP.get(bg_color, [0, 255, 0]), dtype=np.uint8
    )
    a = (alpha_u8[..., None] / 255.0).astype(np.float32)
    return (orig_rgb * a + bg * (1.0 - a)).astype(np.uint8)


def _extract_video_frames(video_path: str, max_frames: int) -> tuple[str, list[str], float]:
    """Extract frames to a temp jpg folder (numeric names for SAM loaders)."""
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise ValueError(f"Cannot open video: {video_path}")

    fps = float(cap.get(cv2.CAP_PROP_FPS) or 25.0)
    if fps <= 1e-3:
        fps = 25.0

    frame_dir = tempfile.mkdtemp(prefix="sam2matting_frames_")
    frame_files: list[str] = []
    idx = 0
    while idx < max_frames:
        ok, frame_bgr = cap.read()
        if not ok:
            break
        name = f"{idx:05d}.jpg"
        out_path = os.path.join(frame_dir, name)
        cv2.imwrite(out_path, frame_bgr, [int(cv2.IMWRITE_JPEG_QUALITY), 95])
        frame_files.append(name)
        idx += 1
    cap.release()

    if not frame_files:
        shutil.rmtree(frame_dir, ignore_errors=True)
        raise ValueError("No frames extracted from the video.")
    return frame_dir, frame_files, fps


def _write_mp4(path: str, frames_bgr: list[np.ndarray], fps: float):
    h, w = frames_bgr[0].shape[:2]
    writer = cv2.VideoWriter(
        path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h)
    )
    for f in frames_bgr:
        writer.write(f)
    writer.release()


@spaces.GPU(duration=90)
def image_matting(
    image: Image.Image,
    mask: Image.Image,
    bg_color: str,
    variant: str,
) -> tuple[Image.Image, Image.Image, Image.Image]:
    """Image matting with SAM2 or SAM3 backbone."""
    if image is None:
        raise ValueError("An input image is required.")
    if variant not in VARIANTS:
        raise ValueError(f"Unknown variant: {variant}")

    image = image.convert("RGB")
    if mask is None:
        mask = _auto_generate_mask(image)
    mask = mask.convert("L")
    if mask.size != image.size:
        mask = mask.resize(image.size, Image.BILINEAR)

    predictor, mask_size = get_image_predictor(variant)
    family = VARIANTS[variant][0]

    with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
        img = predictor.set_image(image)
        raw_mask, mask_input = _prepare_image_mask_tensors(mask, mask_size)

        if family == "sam2":
            _, alpha, _ = predictor.predict(
                img=img,
                raw_mask=raw_mask,
                mask_input=mask_input,
                multimask_output=False,
            )
        else:
            _, alpha, _, _ = predictor.predict(
                img=img,
                raw_mask=raw_mask,
                mask_input=mask_input,
                multimask_output=False,
            )

    alpha_result = (np.asarray(alpha) * 255).astype(np.uint8).squeeze()
    alpha_image = Image.fromarray(alpha_result, mode="L")

    orig_array = np.array(image)
    composite_image = Image.fromarray(_compose_rgb(orig_array, alpha_result, bg_color))
    cutout_image = Image.fromarray(
        np.dstack([orig_array, alpha_result]).astype(np.uint8), mode="RGBA"
    )
    return alpha_image, composite_image, cutout_image


@spaces.GPU(duration=180)
def video_matting(
    video_path: str,
    mask: Image.Image,
    bg_color: str,
    variant: str,
    max_frames: int = 60,
) -> tuple[str, str]:
    """Video matting with SAM2 or SAM3 backbone. Returns (alpha_mp4, composite_mp4)."""
    if video_path is None:
        raise ValueError("An input video is required.")
    if variant not in VARIANTS:
        raise ValueError(f"Unknown variant: {variant}")

    max_frames = int(max_frames)
    frame_dir, frame_files, fps = _extract_video_frames(video_path, max_frames)

    try:
        first_frame = Image.open(
            os.path.join(frame_dir, frame_files[0])
        ).convert("RGB")

        if mask is None:
            mask = _auto_generate_mask(first_frame)
        mask = mask.convert("L")
        if mask.size != first_frame.size:
            mask = mask.resize(first_frame.size, Image.BILINEAR)

        predictor, mask_size = get_video_predictor(variant)
        family = VARIANTS[variant][0]
        device = DEVICE

        # Soft mask logits, same as UniMatting inference_*_video_*.py
        m = np.array(mask).astype(np.float32) / 255.0
        m = (m > 0.005).astype(np.float32) * 20 - 10
        m = torch.from_numpy(m)[None, None]
        m = F.interpolate(
            m, size=(mask_size, mask_size), mode="bilinear", align_corners=False
        )

        alpha_frames_bgr: list[np.ndarray] = []
        comp_frames_bgr: list[np.ndarray] = []

        with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
            inference_state = predictor.init_state(video_path=frame_dir)
            predictor.reset_state(inference_state)
            predictor.add_new_mask(
                inference_state=inference_state,
                frame_idx=0,
                obj_id=1,
                mask=m.to(device),
            )

            for out in predictor.propagate_in_video(inference_state):
                if family == "sam2":
                    out_frame_idx, _, _, alpha, _ = out
                    alpha_2d = (
                        alpha.detach().cpu().squeeze().float().numpy().clip(0, 1)
                    )
                else:
                    out_frame_idx, _, _, alpha, _ = out
                    alpha_2d = np.asarray(alpha).squeeze().clip(0, 1)

                alpha_u8 = (alpha_2d * 255).astype(np.uint8)
                alpha_frames_bgr.append(
                    cv2.cvtColor(alpha_u8, cv2.COLOR_GRAY2BGR)
                )

                orig = np.array(
                    Image.open(
                        os.path.join(frame_dir, frame_files[out_frame_idx])
                    ).convert("RGB")
                )
                comp = _compose_rgb(orig, alpha_u8, bg_color)
                comp_frames_bgr.append(cv2.cvtColor(comp, cv2.COLOR_RGB2BGR))

        out_dir = tempfile.mkdtemp(prefix="sam2matting_out_")
        alpha_path = os.path.join(out_dir, "pha.mp4")
        comp_path = os.path.join(out_dir, "fgr.mp4")
        _write_mp4(alpha_path, alpha_frames_bgr, fps)
        _write_mp4(comp_path, comp_frames_bgr, fps)
        return alpha_path, comp_path
    finally:
        shutil.rmtree(frame_dir, ignore_errors=True)


# --- UI ---

CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            # SAM2Matting: Generalized Image and Video Matting

            Choose a backbone (**SAM2.1-Tiny / SAM2.1-Base+ / SAM3**) and run
            **image** or **video** matting. Optionally provide a rough foreground
            mask; if omitted, one is auto-generated via
            [Inspyrenet-Rembg](https://huggingface.co/spaces/gokaygokay/Inspyrenet-Rembg).

            [Paper](https://arxiv.org/abs/2606.27339) |
            [GitHub](https://github.com/FudanCVL/SAM2Matting) |
            [Model](https://huggingface.co/FudanCVL/SAM2Matting)
            """
        )

        variant = gr.Dropdown(
            label="Backbone",
            choices=list(VARIANTS.keys()),
            value="SAM2.1-Tiny",
        )

        with gr.Tabs():
            with gr.Tab("Image"):
                with gr.Row():
                    with gr.Column(scale=1):
                        input_image = gr.Image(
                            label="Input Image", type="pil", height=400
                        )
                        with gr.Accordion("Foreground Mask (optional)", open=False):
                            input_mask = gr.Image(
                                label="Foreground Mask (white = foreground). "
                                "Leave empty to auto-generate.",
                                type="pil",
                                height=320,
                                image_mode="L",
                            )
                        img_bg = gr.Dropdown(
                            label="Background Color for Composite",
                            choices=list(COLOR_MAP.keys()),
                            value="Green",
                        )
                        img_btn = gr.Button("Matte Image", variant="primary")
                    with gr.Column(scale=1):
                        alpha_out = gr.Image(
                            label="Alpha Matte", type="pil", height=280
                        )
                        composite_out = gr.Image(
                            label="Composite Preview", type="pil", height=280
                        )
                        cutout_out = gr.Image(
                            label="Transparent PNG (alpha applied)",
                            type="pil",
                            image_mode="RGBA",
                            format="png",
                            height=280,
                        )

                img_btn.click(
                    fn=image_matting,
                    inputs=[input_image, input_mask, img_bg, variant],
                    outputs=[alpha_out, composite_out, cutout_out],
                    api_name="matte_image",
                )

                gr.Examples(
                    examples=[
                        ["examples/image.jpg", "examples/mask.png", "Green", "SAM2.1-Tiny"],
                    ],
                    inputs=[input_image, input_mask, img_bg, variant],
                    outputs=[alpha_out, composite_out, cutout_out],
                    fn=image_matting,
                    cache_examples=False,
                )

            with gr.Tab("Video"):
                with gr.Row():
                    with gr.Column(scale=1):
                        input_video = gr.Video(label="Input Video (mp4)")
                        with gr.Accordion(
                            "First-frame Foreground Mask (optional)", open=False
                        ):
                            video_mask = gr.Image(
                                label="Mask for the first frame (white = foreground). "
                                "Leave empty to auto-generate from frame 0.",
                                type="pil",
                                height=320,
                                image_mode="L",
                            )
                        vid_bg = gr.Dropdown(
                            label="Background Color for Composite",
                            choices=list(COLOR_MAP.keys()),
                            value="Green",
                        )
                        max_frames = gr.Slider(
                            label="Max frames (ZeroGPU / time limit)",
                            minimum=8,
                            maximum=150,
                            value=60,
                            step=1,
                        )
                        vid_btn = gr.Button("Matte Video", variant="primary")
                    with gr.Column(scale=1):
                        alpha_video_out = gr.Video(label="Alpha Matte Video")
                        composite_video_out = gr.Video(
                            label="Composite Preview Video"
                        )

                vid_btn.click(
                    fn=video_matting,
                    inputs=[input_video, video_mask, vid_bg, variant, max_frames],
                    outputs=[alpha_video_out, composite_video_out],
                    api_name="matte_video",
                )
                
                gr.Examples(
                    examples=[
                        ["examples/demo_video.mp4", "examples/video_mask.png", "Green", "SAM2.1-Tiny", 30],
                    ],
                    inputs=[input_video, video_mask, vid_bg, variant, max_frames],
                    outputs=[alpha_video_out, composite_video_out],
                    fn=video_matting,
                    cache_examples=False,
                )

        gr.Markdown(
            """
            ### Tips
            - Mask should roughly cover the foreground (white = foreground).
            - Soft grayscale masks work best.
            - **SAM2.1-Tiny** is fastest; **SAM3** is heaviest (needs more VRAM/time).
            - Video matting uses the first-frame mask and propagates through the clip.
            - Keep `Max frames` modest on ZeroGPU Spaces.

            ### License
            CC-BY-NC-SA-4.0 (non-commercial research use only).
            """
        )

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