import gradio as gr
import cv2
import tempfile
import os
import numpy as np
import spaces
from ultralytics import YOLO

# ── Example video (place video.mp4 in the repo root) ───────────────────────
EXAMPLE_VIDEO_PATH = "video.mp4"

# ── COCO 80 classes for reference hint ─────────────────────────────────────
COCO_CLASSES = [
    "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train",
    "truck", "boat", "traffic light", "fire hydrant", "stop sign",
    "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
    "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag",
    "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball", "kite",
    "baseball bat", "baseball glove", "skateboard", "surfboard",
    "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon",
    "bowl", "banana", "apple", "sandwich", "orange", "broccoli", "carrot",
    "hot dog", "pizza", "donut", "cake", "chair", "couch", "potted plant",
    "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote",
    "keyboard", "cell phone", "microwave", "oven", "toaster", "sink",
    "refrigerator", "book", "clock", "vase", "scissors", "teddy bear",
    "hair drier", "toothbrush",
]

# ── Model cache ─────────────────────────────────────────────────────────────
model = None


def _parse_class_filter(class_prompt: str, model_names: dict):
    """
    Parse a comma-separated prompt like "person, car, dog" into
    (list[int] | None, list[str]).
    Returns (None, []) to detect all classes when prompt is empty.
    """
    if not class_prompt or not class_prompt.strip():
        return None, []

    name_to_id = {v.lower(): k for k, v in model_names.items()}
    requested  = [t.strip().lower() for t in class_prompt.split(",") if t.strip()]
    matched, unknown = [], []

    for name in requested:
        if name in name_to_id:
            matched.append(name_to_id[name])
        else:
            unknown.append(name)

    return (matched if matched else None), unknown


# ── Core inference ──────────────────────────────────────────────────────────
@spaces.GPU(duration=180)
def process_video(
    video_path: str,
    conf_threshold: float,
    iou_threshold: float,
    model_size: str,
    task: str,
    class_prompt: str,
) -> tuple:
    """Run YOLO26 inference on an uploaded video and return annotated video + stats."""

    if video_path is None:
        return None, "⚠️ 請先上傳影片"

    global model
    model_map = {
        "Nano (最快)":   "yolo26n.pt",
        "Small":         "yolo26s.pt",
        "Medium":        "yolo26m.pt",
        "Large":         "yolo26l.pt",
        "XLarge (最準)": "yolo26x.pt",
    }
    task_map = {
        "物件偵測 (detect)":  "detect",
        "實例分割 (segment)": "segment",
        "姿態估計 (pose)":    "pose",
    }

    selected_pt   = model_map.get(model_size, "yolo26n.pt")
    selected_task = task_map.get(task, "detect")

    if selected_task == "detect":
        pt_name = selected_pt
    else:
        suffix  = "-seg" if selected_task == "segment" else "-pose"
        pt_name = selected_pt.replace(".pt", f"{suffix}.pt")

    if model is None or getattr(model, "ckpt_path", None) != pt_name:
        try:
            model = YOLO(pt_name)
        except Exception as e:
            model   = YOLO("yolo26n.pt")
            pt_name = "yolo26n.pt"
            return None, f"⚠️ 無法載入 {pt_name}，已回退至 yolo26n.pt\n錯誤：{e}"

    # ── Parse class filter ──────────────────────────────────────────────────
    class_ids, unknown_names = _parse_class_filter(class_prompt, model.names)

    # ── Open video ──────────────────────────────────────────────────────────
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        return None, "❌ 無法開啟影片檔案"

    fps    = cap.get(cv2.CAP_PROP_FPS) or 25
    width  = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    out_path = tempfile.mktemp(suffix=".mp4")
    fourcc   = cv2.VideoWriter_fourcc(*"mp4v")
    writer   = cv2.VideoWriter(out_path, fourcc, fps, (width, height))

    frame_idx        = 0
    total_detections = 0
    class_counts     = {}

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        infer_kwargs = dict(conf=conf_threshold, iou=iou_threshold, verbose=False)
        if class_ids:
            infer_kwargs["classes"] = class_ids   # YOLO native class-filter param

        results   = model(frame, **infer_kwargs)
        annotated = results[0].plot()

        boxes = results[0].boxes
        if boxes is not None:
            total_detections += len(boxes)
            for cls_id in boxes.cls.cpu().numpy().astype(int):
                name = model.names[cls_id]
                class_counts[name] = class_counts.get(name, 0) + 1

        writer.write(annotated)
        frame_idx += 1

    cap.release()
    writer.release()

    # ── Build stats ─────────────────────────────────────────────────────────
    if class_ids:
        filter_desc = "🎯 過濾類別：" + ", ".join(model.names[i] for i in class_ids)
    else:
        filter_desc = "🎯 過濾類別：全部（無限制）"

    stats_lines = [
        "✅ 處理完成",
        f"📽️  總幀數：{frame_idx}",
        f"🔍 總偵測次數：{total_detections}",
        f"⚡ 模型：{pt_name}",
        filter_desc,
    ]
    if unknown_names:
        stats_lines.append(f"⚠️ 找不到類別：{', '.join(unknown_names)}（已忽略）")

    stats_lines += ["", "📊 各類別累計偵測次數："]
    if class_counts:
        for cls, cnt in sorted(class_counts.items(), key=lambda x: -x[1]):
            stats_lines.append(f"  • {cls}: {cnt}")
    else:
        stats_lines.append("  （未偵測到任何物件）")

    # ── Re-encode to H.264 ──────────────────────────────────────────────────
    final_path = tempfile.mktemp(suffix=".mp4")
    os.system(
        f'ffmpeg -y -i "{out_path}" -vcodec libx264 -acodec aac "{final_path}" -loglevel quiet'
    )
    if not os.path.exists(final_path):
        final_path = out_path

    return final_path, "\n".join(stats_lines)


# ── Gradio UI ───────────────────────────────────────────────────────────────
CSS = """
@import url('https://fonts.googleapis.com/css2?family=Space+Mono:wght@400;700&family=Syne:wght@400;700;800&display=swap');

:root {
    --bg:      #0a0a0f;
    --surface: #111118;
    --border:  #1e1e2e;
    --accent:  #00ff88;
    --accent2: #7c3aed;
    --text:    #e2e8f0;
    --muted:   #64748b;
    --warn:    #f59e0b;
    --radius:  12px;
}
body, .gradio-container {
    background: var(--bg) !important;
    color: var(--text) !important;
    font-family: 'Syne', sans-serif !important;
}
h1, h2, h3 { font-family: 'Syne', sans-serif !important; }

#header {
    text-align: center;
    padding: 2.5rem 1rem 1.5rem;
    background: linear-gradient(135deg, #0a0a0f 0%, #111118 100%);
    border-bottom: 1px solid var(--border);
}
#header h1 {
    font-size: 2.8rem;
    font-weight: 800;
    background: linear-gradient(90deg, var(--accent) 0%, #00ccff 50%, var(--accent2) 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    margin: 0 0 0.4rem;
    letter-spacing: -1px;
}
#header p { color: var(--muted); font-family: 'Space Mono', monospace; font-size: 0.85rem; }

.gr-group, .gr-box, .gr-form {
    background: var(--surface) !important;
    border: 1px solid var(--border) !important;
    border-radius: var(--radius) !important;
}
input[type=range]::-webkit-slider-thumb { background: var(--accent) !important; }

button.primary {
    background: linear-gradient(135deg, var(--accent) 0%, #00ccff 100%) !important;
    color: #000 !important; font-weight: 700 !important;
    font-family: 'Syne', sans-serif !important;
    border: none !important; border-radius: 8px !important;
    letter-spacing: 0.5px;
}
button.secondary {
    background: transparent !important;
    border: 1px solid var(--border) !important;
    color: var(--text) !important;
}
textarea, .gr-textbox {
    background: var(--bg) !important;
    color: var(--accent) !important;
    font-family: 'Space Mono', monospace !important;
    font-size: 0.82rem !important;
    border: 1px solid var(--border) !important;
}
#class-filter textarea {
    color: var(--warn) !important;
    border-color: rgba(245,158,11,0.5) !important;
}
label span { color: var(--muted) !important; font-size: 0.8rem; }
.tag {
    display: inline-block;
    background: rgba(0,255,136,0.1); color: var(--accent);
    border: 1px solid rgba(0,255,136,0.3);
    padding: 2px 10px; border-radius: 999px;
    font-size: 0.75rem; font-family: 'Space Mono', monospace; margin: 2px;
}
.tag-cls {
    background: rgba(245,158,11,0.12); color: var(--warn);
    border-color: rgba(245,158,11,0.35);
}
"""

HEADER_HTML = """
<div id="header">
  <h1>⚡ YOLO26 Video Detector</h1>
  <p>Ultralytics YOLO26 · End-to-End NMS-Free · Real-time Object Detection</p>
  <div style="margin-top:0.8rem">
    <span class="tag">detect</span>
    <span class="tag">segment</span>
    <span class="tag">pose</span>
    <span class="tag">NMS-free</span>
    <span class="tag tag-cls">🎯 class filter</span>
  </div>
</div>
"""

COCO_HINT_HTML = (
    "<details style='margin-top:6px;font-family:monospace;font-size:0.75rem;color:#64748b'>"
    "<summary style='cursor:pointer;color:#f59e0b;user-select:none'>"
    "▶ 可用 COCO 類別清單（點擊展開）</summary>"
    "<div style='margin-top:8px;line-height:2;'>"
    + "  ".join(f"<code style='background:#1e1e2e;padding:1px 5px;border-radius:4px'>{c}</code>"
                for c in COCO_CLASSES)
    + "</div></details>"
)

with gr.Blocks(css=CSS, title="YOLO26 Video Detector") as demo:
    gr.HTML(HEADER_HTML)

    with gr.Row(equal_height=False):
        # ── Left column ──────────────────────────────────────────────────────
        with gr.Column(scale=1):
            gr.Markdown("### 📤 上傳影片")
            video_input = gr.Video(label="輸入影片", sources=["upload"])

            gr.Markdown("### ⚙️ 推論設定")
            with gr.Group():
                model_size = gr.Dropdown(
                    label="模型大小",
                    choices=["Nano (最快)", "Small", "Medium", "Large", "XLarge (最準)"],
                    value="Nano (最快)",
                )
                task = gr.Dropdown(
                    label="任務類型",
                    choices=["物件偵測 (detect)", "實例分割 (segment)", "姿態估計 (pose)"],
                    value="物件偵測 (detect)",
                )
                conf = gr.Slider(0.1, 1.0, step=0.05, value=0.25, label="Confidence 閾值")
                iou  = gr.Slider(0.1, 1.0, step=0.05, value=0.45, label="IoU 閾值")

            gr.Markdown("### 🎯 類別過濾 Prompt")
            with gr.Group():
                class_filter = gr.Textbox(
                    elem_id="class-filter",
                    label="指定偵測類別（留空 = 偵測全部）",
                    placeholder="例：person, car, dog",
                    lines=1,
                )
                gr.HTML(COCO_HINT_HTML)

            with gr.Row():
                run_btn   = gr.Button("🚀 開始偵測", variant="primary")
                clear_btn = gr.Button("🗑️ 清除", variant="secondary")

        # ── Right column ─────────────────────────────────────────────────────
        with gr.Column(scale=1):
            gr.Markdown("### 📥 偵測結果")
            video_output = gr.Video(label="輸出影片", interactive=False)
            stats_output = gr.Textbox(
                label="統計資訊", lines=14, interactive=False,
                placeholder="偵測完成後，統計資訊將顯示於此…",
            )

    # ── Examples ─────────────────────────────────────────────────────────────
    gr.Markdown("---")
    gr.Markdown("#### 🎬 範例影片")

    gr.Examples(
        examples=[
            [EXAMPLE_VIDEO_PATH, 0.25, 0.45, "Nano (最快)", "物件偵測 (detect)",  ""],
            [EXAMPLE_VIDEO_PATH, 0.25, 0.45, "Nano (最快)", "物件偵測 (detect)",  "person, car"],
            [EXAMPLE_VIDEO_PATH, 0.25, 0.45, "Nano (最快)", "物件偵測 (detect)",  "dog, cat, bird"],
            [EXAMPLE_VIDEO_PATH, 0.40, 0.45, "Nano (最快)", "實例分割 (segment)", "person"],
            [EXAMPLE_VIDEO_PATH, 0.25, 0.45, "Small",       "姿態估計 (pose)",    "person"],
        ],
        inputs=[video_input, conf, iou, model_size, task, class_filter],
        label="點擊範例自動填入設定，再按「🚀 開始偵測」",
        examples_per_page=5,
    )

    gr.Markdown("#### 💡 使用提示")
    gr.Markdown(
        "- **類別過濾**：輸入英文類別名（逗號分隔），例如 `person, car`；留空則偵測全部\n"
        "- 類別名稱不區分大小寫，拼錯的類別會在統計欄位顯示警告並忽略\n"
        "- **Nano** 模型最快，適合 CPU / 輕量 GPU\n"
        "- **Confidence** 調低可偵測更多物件（但可能有雜訊）\n"
        "- 影片長度建議 < 60 秒以避免逾時\n"
        "- 分割與姿態任務需對應的 `-seg` / `-pose` 權重，初次使用會自動下載"
    )

    # ── Events ───────────────────────────────────────────────────────────────
    run_btn.click(
        fn=process_video,
        inputs=[video_input, conf, iou, model_size, task, class_filter],
        outputs=[video_output, stats_output],
    )

    def clear_all():
        return None, None, "", ""

    clear_btn.click(
        fn=clear_all,
        outputs=[video_input, video_output, stats_output, class_filter],
    )


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