"""
RT-DETR Object Detection — streaming variant (Plan B).

Use this ONLY if you specifically need progressive playback while processing.
It is inherently more fragile than app.py (single-file). Differences from upstream:

  * Every yielded segment is re-encoded to real H.264 with ffmpeg. Gradio's
    streaming-outputs guide requires .mp4 or h.264-in-.ts for video chunks.
  * Segment length is time-based and >= 1 s (Gradio requires this for smooth
    playback). Upstream hardcoded exactly 2*desired_fps frames, so any video
    whose frame count was not an exact multiple of that lost its tail.
  * The trailing partial segment is flushed instead of discarded.
  * Even dimensions enforced; frames resized to the writer's exact size.
"""

import os
import time
import uuid
import shutil
import subprocess

import cv2
import numpy as np
import torch
import gradio as gr
import spaces
from PIL import Image
from transformers import RTDetrForObjectDetection, RTDetrImageProcessor

from draw_boxes import draw_bounding_boxes

BUILD_TAG = "PLANB-H264-STREAMING-v1"
print("=" * 70, flush=True)
print(f"[BOOT] RT-DETR build tag: {BUILD_TAG}", flush=True)
print(f"[BOOT] ffmpeg on PATH: {shutil.which('ffmpeg')}", flush=True)
print("=" * 70, flush=True)

MODEL_ID = "PekingU/rtdetr_r50vd"
image_processor = RTDetrImageProcessor.from_pretrained(MODEL_ID)
model = RTDetrForObjectDetection.from_pretrained(MODEL_ID).to("cuda")

SUBSAMPLE = 2
SEGMENT_SECONDS = 1.5          # must stay >= 1.0
WORK_DIR = os.path.join(os.getcwd(), "rtdetr_stream")
os.makedirs(WORK_DIR, exist_ok=True)


def encode_h264(src_path: str) -> str:
    dst_path = src_path.replace(".mp4", "_h264.mp4")
    cmd = [
        "ffmpeg", "-y", "-i", src_path,
        "-c:v", "libx264", "-profile:v", "baseline", "-level", "3.1",
        "-preset", "veryfast", "-crf", "23",
        "-pix_fmt", "yuv420p", "-movflags", "+faststart", "-an",
        dst_path,
    ]
    proc = subprocess.run(cmd, capture_output=True, text=True)
    if proc.returncode != 0 or not os.path.exists(dst_path):
        print("[H264] FAILED:", proc.stderr[-800:], flush=True)
        raise gr.Error("H.264 re-encode failed — see container logs.")
    os.remove(src_path)
    print(f"[H264] OK -> {dst_path} ({os.path.getsize(dst_path)} bytes)", flush=True)
    return dst_path


@spaces.GPU(duration=180)
def stream_object_detection(video, conf_threshold):
    print(f"\n[RUN] {BUILD_TAG} | conf={conf_threshold}", flush=True)
    cap = cv2.VideoCapture(video)
    if not cap.isOpened():
        raise gr.Error("Could not open the uploaded video.")

    src_fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
    out_fps = max(1.0, src_fps / SUBSAMPLE)
    width = max(2, ((int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) // 2) // 2) * 2)
    height = max(2, ((int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) // 2) // 2) * 2)
    per_segment = max(1, int(round(out_fps * SEGMENT_SECONDS)))
    print(f"[RUN] out {width}x{height} @ {out_fps:.2f}fps | {per_segment} frames/segment",
          flush=True)

    def new_writer():
        p = os.path.join(WORK_DIR, f"seg_{uuid.uuid4().hex}.mp4")
        return cv2.VideoWriter(p, cv2.VideoWriter_fourcc(*"mp4v"), out_fps, (width, height)), p

    writer, path = new_writer()
    batch, n_read, seg_no = [], 0, 0

    def flush(frames, wr, p):
        if not frames:
            wr.release()
            if os.path.exists(p):
                os.remove(p)
            return None
        inputs = image_processor(images=frames, return_tensors="pt").to("cuda")
        t0 = time.time()
        with torch.no_grad():
            outputs = model(**inputs)
        boxes = image_processor.post_process_object_detection(
            outputs,
            target_sizes=torch.tensor([(height, width)] * len(frames)),
            threshold=conf_threshold,
        )
        n_det = 0
        for array, box in zip(frames, boxes):
            n_det += int(box["scores"].shape[0])
            pil_image = draw_bounding_boxes(Image.fromarray(array), box, model, conf_threshold)
            wr.write(np.array(pil_image)[:, :, ::-1].copy())
        wr.release()
        print(f"[SEG] frames={len(frames)} inference={time.time() - t0:.3f}s dets={n_det}",
              flush=True)
        return encode_h264(p)

    try:
        while True:
            ok, frame = cap.read()
            if not ok:
                break
            if n_read % SUBSAMPLE == 0:
                frame = cv2.resize(frame, (width, height))
                batch.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
                if len(batch) >= per_segment:
                    out = flush(batch, writer, path)
                    if out:
                        seg_no += 1
                        print(f"[YIELD] #{seg_no} {out}", flush=True)
                        yield out
                    batch = []
                    writer, path = new_writer()
            n_read += 1

        # Trailing frames — upstream dropped these.
        out = flush(batch, writer, path)
        if out:
            seg_no += 1
            print(f"[YIELD] #{seg_no} (tail) {out}", flush=True)
            yield out
    finally:
        cap.release()

    print(f"[RUN] done. read={n_read} segments={seg_no}", flush=True)
    if seg_no == 0:
        raise gr.Error("No segments produced — the video decoded to zero frames.")


with gr.Blocks(title="RT-DETR Streaming") as app:
    gr.HTML(
        "<h1 style='text-align:center;margin-bottom:0'>RT-DETR Object Detection (streaming)</h1>"
        f"<p style='text-align:center;color:#888;font-size:0.85em'>build {BUILD_TAG}</p>"
    )
    with gr.Row():
        with gr.Column():
            video = gr.Video(label="Video Source")
            conf_threshold = gr.Slider(
                label="Confidence Threshold",
                minimum=0.0, maximum=1.0, step=0.01, value=0.30,
            )
        with gr.Column():
            output_video = gr.Video(label="Processed Video", streaming=True, autoplay=True)

    video.upload(
        fn=stream_object_detection,
        inputs=[video, conf_threshold],
        outputs=[output_video],
    )

if __name__ == "__main__":
    app.launch()