import tempfile
from pathlib import Path

import gradio as gr
import numpy as np
import spaces  # noqa: F401
import supervision as sv
import torch
from PIL import Image

from transformers import (
    AutoImageProcessor,
    AutoModelForInstanceSegmentation,
    AutoModelForObjectDetection,
)


WEBCAM_MAX_SHORT_EDGE = 384


def _get_device() -> str:
    if torch.cuda.is_available():
        return "cuda"
    if torch.backends.mps.is_available():
        return "mps"
    return "cpu"


DEVICE = _get_device()
DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32

if DEVICE == "cuda":
    torch.backends.cudnn.benchmark = True

BOX_ANNOTATOR = sv.BoxAnnotator(thickness=2)
LABEL_ANNOTATOR = sv.LabelAnnotator(text_scale=0.6, text_thickness=1, text_padding=4)
MASK_ANNOTATOR = sv.MaskAnnotator(opacity=0.35)

gr.Warning("Loading models — this may take a moment on first launch...")

DET_PROCESSOR = AutoImageProcessor.from_pretrained("Roboflow/rf-detr-medium", revision="refs/pr/1", device=DEVICE)
DET_MODEL = AutoModelForObjectDetection.from_pretrained(
    "Roboflow/rf-detr-medium", torch_dtype=DTYPE, attn_implementation="sdpa",
).to(DEVICE).eval()

WEBCAM_SEG_PROCESSOR = AutoImageProcessor.from_pretrained("Roboflow/rf-detr-seg-medium", revision="refs/pr/1", device=DEVICE)
WEBCAM_SEG_MODEL = AutoModelForInstanceSegmentation.from_pretrained(
    "Roboflow/rf-detr-seg-medium", torch_dtype=DTYPE, attn_implementation="sdpa",
).to(DEVICE).eval()

UPLOAD_SEG_PROCESSOR = AutoImageProcessor.from_pretrained("Roboflow/rf-detr-seg-medium", device=DEVICE)
UPLOAD_SEG_MODEL = AutoModelForInstanceSegmentation.from_pretrained(
    "Roboflow/rf-detr-seg-medium", torch_dtype=DTYPE, attn_implementation="sdpa",
).to(DEVICE).eval()

DET_ID2LABEL = DET_MODEL.config.id2label
WEBCAM_SEG_ID2LABEL = WEBCAM_SEG_MODEL.config.id2label
UPLOAD_SEG_ID2LABEL = UPLOAD_SEG_MODEL.config.id2label


def _downscale(image_array: np.ndarray, max_short_edge: int) -> np.ndarray:
    h, w = image_array.shape[:2]
    short_edge = min(h, w)
    if short_edge <= max_short_edge:
        return image_array
    scale = max_short_edge / short_edge
    new_w, new_h = int(w * scale), int(h * scale)
    return np.array(Image.fromarray(image_array).resize((new_w, new_h), Image.BILINEAR))


def _run_detection(image_array: np.ndarray, score_threshold: float) -> np.ndarray:
    inputs = DET_PROCESSOR(images=image_array, return_tensors="pt").to(DEVICE)

    with torch.inference_mode():
        outputs = DET_MODEL(**inputs)

    h, w = image_array.shape[:2]
    results = DET_PROCESSOR.post_process_object_detection(
        outputs, threshold=score_threshold, target_sizes=torch.tensor([[h, w]], device=DEVICE)
    )[0]

    detections = sv.Detections.from_transformers(transformers_results=results, id2label=DET_ID2LABEL)
    if len(detections) == 0:
        return image_array

    labels = [
        f"{DET_ID2LABEL[cid]} {c:.2f}"
        for cid, c in zip(detections.class_id, detections.confidence)
    ]
    scene = BOX_ANNOTATOR.annotate(scene=image_array.copy(), detections=detections)
    return LABEL_ANNOTATOR.annotate(scene=scene, detections=detections, labels=labels)


def _run_segmentation(
    image_array: np.ndarray,
    processor,
    model,
    id2label: dict,
    score_threshold: float,
    mask_threshold: float,
) -> np.ndarray:
    inputs = processor(images=image_array, return_tensors="pt").to(DEVICE)

    with torch.inference_mode():
        outputs = model(**inputs)

    h, w = image_array.shape[:2]
    result = processor.post_process_instance_segmentation(
        outputs, threshold=score_threshold, mask_threshold=mask_threshold, target_sizes=[(h, w)],
    )[0]

    segments = result["segments_info"]
    if not segments:
        return image_array

    segmentation_map = result["segmentation"].cpu().numpy()
    n = len(segments)
    masks = np.empty((n, h, w), dtype=bool)
    confidences = np.empty(n, dtype=np.float32)
    class_ids = np.empty(n, dtype=int)
    boxes = np.empty((n, 4), dtype=np.float32)

    keep = 0
    for seg in segments:
        binary_mask = segmentation_map == seg["id"]
        ys, xs = np.where(binary_mask)
        if ys.size == 0:
            continue
        masks[keep] = binary_mask
        confidences[keep] = seg["score"]
        class_ids[keep] = seg["label_id"]
        boxes[keep] = [xs.min(), ys.min(), xs.max(), ys.max()]
        keep += 1

    if keep == 0:
        return image_array

    detections = sv.Detections(
        xyxy=boxes[:keep], mask=masks[:keep],
        confidence=confidences[:keep], class_id=class_ids[:keep],
    )
    labels = [
        f"{id2label[cid]} {c:.2f}"
        for cid, c in zip(detections.class_id, detections.confidence)
    ]
    scene = MASK_ANNOTATOR.annotate(scene=image_array.copy(), detections=detections)
    scene = BOX_ANNOTATOR.annotate(scene=scene, detections=detections)
    return LABEL_ANNOTATOR.annotate(scene=scene, detections=detections, labels=labels)


def infer_webcam_frame(
    image_array: np.ndarray,
    task: str,
    score_threshold: float,
    mask_threshold: float,
) -> np.ndarray:
    image_array = _downscale(image_array, WEBCAM_MAX_SHORT_EDGE)
    if task == "Object detection":
        return _run_detection(image_array, score_threshold)
    return _run_segmentation(
        image_array, WEBCAM_SEG_PROCESSOR, WEBCAM_SEG_MODEL, WEBCAM_SEG_ID2LABEL,
        score_threshold, mask_threshold,
    )


def infer_single_image(
    image_array: np.ndarray,
    task: str,
    score_threshold: float,
    mask_threshold: float,
) -> np.ndarray:
    if task == "Object detection":
        return _run_detection(image_array, score_threshold)
    return _run_segmentation(
        image_array, UPLOAD_SEG_PROCESSOR, UPLOAD_SEG_MODEL, UPLOAD_SEG_ID2LABEL,
        score_threshold, mask_threshold,
    )


def process_video(
    video_path: str,
    task: str,
    score_threshold: float,
    mask_threshold: float,
) -> str:
    import cv2

    if not video_path:
        raise gr.Error("Please upload a video.")

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

    fps = 30
    width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    output_path = str(Path(tempfile.mkdtemp(prefix="rf-detr-demo-")) / "annotated.mp4")
    writer = cv2.VideoWriter(output_path, fourcc, fps, (width, height))

    while True:
        success, frame = capture.read()
        if not success:
            break
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        processed_rgb = infer_single_image(frame_rgb, task, score_threshold, mask_threshold)
        writer.write(cv2.cvtColor(processed_rgb, cv2.COLOR_RGB2BGR))

    capture.release()
    writer.release()
    return output_path


with gr.Blocks(title="RF-DETR Realtime Webcam + Upload") as demo:
    gr.Markdown(
        "# RF-DETR Realtime Demo\n"
        "Live webcam + upload tabs with object detection and instance segmentation.\n"
        "Uses Roboflow RF-DETR checkpoints."
    )

    with gr.Tab("Realtime Webcam"):
        gr.Markdown(
            "Uses **rf-detr-base-2** for detection and **rf-detr-seg-medium** for segmentation "
            "(optimized for real-time speed)."
        )
        webcam_stream = gr.Image(
            label="Webcam stream",
            sources="webcam",
            type="numpy",
            streaming=True,
        )
        webcam_task = gr.Radio(
            ["Object detection", "Instance segmentation"],
            value="Object detection",
            label="Task",
        )
        webcam_score_threshold = gr.Slider(0.0, 1.0, value=0.4, step=0.01, label="Score threshold")
        webcam_mask_threshold = gr.Slider(0.0, 1.0, value=0.5, step=0.01, label="Mask threshold")

        webcam_stream.stream(
            fn=infer_webcam_frame,
            inputs=[
                webcam_stream,
                webcam_task,
                webcam_score_threshold,
                webcam_mask_threshold,
            ],
            outputs=webcam_stream,
            queue=False,
            concurrency_limit=1,
            trigger_mode="always_last",
            show_progress="hidden",
            stream_every=0.005,
        )

    with gr.Tab("Upload Images & Videos"):
        gr.Markdown(
            "Uses **rf-detr-base-2** for detection and **rf-detr-seg-xxlarge** for segmentation "
            "(maximum quality)."
        )
        upload_task = gr.Radio(
            ["Object detection", "Instance segmentation"],
            value="Object detection",
            label="Task",
        )
        upload_score_threshold = gr.Slider(0.0, 1.0, value=0.4, step=0.01, label="Score threshold")
        upload_mask_threshold = gr.Slider(0.0, 1.0, value=0.5, step=0.01, label="Mask threshold")

        with gr.Row():
            image_input = gr.Image(label="Upload image", sources=["upload"], type="numpy")
            image_output = gr.Image(label="Annotated image", type="numpy")
        image_button = gr.Button("Run on image", variant="primary")
        image_button.click(
            fn=infer_single_image,
            inputs=[image_input, upload_task, upload_score_threshold, upload_mask_threshold],
            outputs=image_output,
        )

        with gr.Row():
            video_input = gr.Video(label="Upload video", sources=["upload"])
            video_output = gr.Video(label="Annotated video")
        video_button = gr.Button("Run on video", variant="primary")
        video_button.click(
            fn=process_video,
            inputs=[video_input, upload_task, upload_score_threshold, upload_mask_threshold],
            outputs=video_output,
        )


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=2).launch()