import os
import shutil
import subprocess
import tempfile
from functools import lru_cache
from pathlib import Path

import cv2
import gradio as gr
import torch
import spaces
from huggingface_hub import hf_hub_download
from ultralytics.models.sam import SAM3VideoSemanticPredictor


WEIGHTS_REPO_ID = os.getenv("SAM3_WEIGHTS_REPO_ID", "kamillkate/sam3-weights")
WEIGHTS_FILENAME = os.getenv("SAM3_WEIGHTS_FILENAME", "sam3.pt")
WEIGHTS_REPO_TYPE = os.getenv("SAM3_WEIGHTS_REPO_TYPE", "model")


def _video_path(video):
    if video is None:
        raise gr.Error("Please upload a video first.")
    if isinstance(video, dict):
        return video.get("path") or video.get("name")
    return video


def _parse_prompts(prompt_text):
    prompts = [item.strip() for item in prompt_text.split(",") if item.strip()]
    if not prompts:
        raise gr.Error("Enter at least one text prompt, for example: cellphone")
    return prompts


@lru_cache(maxsize=1)
def get_model_path():
    return hf_hub_download(
        repo_id=WEIGHTS_REPO_ID,
        filename=WEIGHTS_FILENAME,
        repo_type=WEIGHTS_REPO_TYPE,
        token=os.getenv("HF_TOKEN"),
    )


@lru_cache(maxsize=8)
def get_predictor(confidence, img_size, use_half):
    model_path = get_model_path()
    overrides = {
        "conf": float(confidence),
        "task": "segment",
        "mode": "predict",
        "imgsz": int(img_size),
        "model": model_path,
        "half": bool(use_half and torch.cuda.is_available()),
        "save": False,
    }
    return SAM3VideoSemanticPredictor(overrides=overrides)


def _to_browser_mp4(source_path):
    ffmpeg = shutil.which("ffmpeg")
    if not ffmpeg:
        return source_path

    encoded_path = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4").name
    cmd = [
        ffmpeg,
        "-y",
        "-i",
        source_path,
        "-c:v",
        "libx264",
        "-pix_fmt",
        "yuv420p",
        "-movflags",
        "+faststart",
        encoded_path,
    ]
    try:
        subprocess.run(cmd, check=True, capture_output=True)
        return encoded_path
    except Exception:
        return source_path

@spaces.GPU(duration=180)
def segment_video(video, prompt_text, confidence, img_size, use_half, progress=gr.Progress()):
    input_path = _video_path(video)
    prompts = _parse_prompts(prompt_text)

    cap = cv2.VideoCapture(input_path)
    if not cap.isOpened():
        raise gr.Error("Could not open the uploaded video.")

    fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) or 1
    cap.release()

    raw_output = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4").name
    writer = cv2.VideoWriter(
        raw_output,
        cv2.VideoWriter_fourcc(*"mp4v"),
        fps,
        (width, height),
    )
    if not writer.isOpened():
        raise gr.Error("Could not create the output video.")

    predictor = get_predictor(round(float(confidence), 3), int(img_size), bool(use_half))
    results = predictor(source=input_path, text=prompts, stream=True)

    frame_count = 0
    detected_count = 0
    try:
        for result in results:
            frame_count += 1
            annotated_frame = result.plot()

            if annotated_frame.shape[1] != width or annotated_frame.shape[0] != height:
                annotated_frame = cv2.resize(annotated_frame, (width, height))

            writer.write(annotated_frame)

            if result.boxes is not None:
                detected_count += len(result.boxes)

            progress(frame_count / total_frames, desc=f"Processing frame {frame_count}/{total_frames}")
    finally:
        writer.release()

    output_path = _to_browser_mp4(raw_output)
    summary = (
        f"Done. Processed {frame_count} frames. "
        f"Total detections across frames: {detected_count}. "
        f"Prompts: {', '.join(prompts)}"
    )
    return output_path, summary


with gr.Blocks(title="SAM 3 Video Concept Segmentation") as demo:
    gr.Markdown(
        """
        # SAM 3 Video Concept Segmentation

        Upload a video, enter one or more text prompts, and generate an annotated segmentation video.
        """
    )

    with gr.Row():
        with gr.Column():
            video_input = gr.Video(label="Input video")
            prompt_input = gr.Textbox(
                value="cellphone",
                label="Text prompts",
                info="Separate multiple concepts with commas, for example: cellphone, hand",
            )
            confidence_input = gr.Slider(
                minimum=0.05,
                maximum=0.95,
                value=0.25,
                step=0.05,
                label="Confidence threshold",
            )
            img_size_input = gr.Dropdown(
                choices=[512, 640, 768, 1024],
                value=640,
                label="Image size",
            )
            half_input = gr.Checkbox(
                value=True,
                label="Use FP16 when CUDA is available",
            )
            run_button = gr.Button("Segment video", variant="primary")

        with gr.Column():
            video_output = gr.Video(label="Output video")
            status_output = gr.Textbox(label="Status", interactive=False)

    inputs = [video_input, prompt_input, confidence_input, img_size_input, half_input]
    run_button.click(
        fn=segment_video,
        inputs=inputs,
        outputs=[video_output, status_output],
        show_progress="full",
    )

    if Path("video.mp4").exists():
        gr.Examples(
            examples=[["video.mp4", "cellphone", 0.25, 640, True]],
            inputs=inputs,
        )


if __name__ == "__main__":
    demo.queue(max_size=4).launch()