"""
SAM 3 視頻概念分割 - 使用文本提示跟踪概念
基於 Ultralytics SAM 3 文檔：
https://docs.ultralytics.com/zh/models/sam-3/#track-concepts-with-text-prompts

需求：
    pip install -U ultralytics
    pip uninstall clip -y
    pip install git+https://github.com/ultralytics/CLIP.git

注意：
    sam3.pt 需要從 Hugging Face 手動下載（需申請權限）：
    https://huggingface.co/facebook/sam3
"""

import cv2
from pathlib import Path
from ultralytics.models.sam import SAM3VideoSemanticPredictor


# ─────────────────────────────────────────────
# 設定區域（請依需求修改）
# ─────────────────────────────────────────────

# 輸入影片路徑（可改為本地影片路徑）
VIDEO_SOURCE = "video.mp4"

# 文本提示：指定要追蹤的概念（可自由修改）
TEXT_PROMPTS = ["cellphone"]

# SAM 3 模型權重路徑（需手動下載）
from huggingface_hub import hf_hub_download
weights_path = hf_hub_download(
    repo_id="kamillkate/sam3-weights",
    filename="sam3.pt",
    repo_type="model",
)
MODEL_PATH = weights_path

# 輸出影片路徑
OUTPUT_PATH = "output.mp4"

# 信心閾值
CONFIDENCE = 0.25

# 圖像尺寸
IMG_SIZE = 640

# ─────────────────────────────────────────────
# 主程式
# ─────────────────────────────────────────────

def main():
    print("=" * 60)
    print("SAM 3 視頻概念分割 - 文本提示跟踪")
    print("=" * 60)
    print(f"📹 輸入影片：{VIDEO_SOURCE}")
    print(f"🔍 追蹤概念：{TEXT_PROMPTS}")
    print(f"💾 輸出路徑：{OUTPUT_PATH}")
    print()

    # 確認影片存在
    if not Path(VIDEO_SOURCE).exists():
        raise FileNotFoundError(
            f"找不到影片：{VIDEO_SOURCE}\n"
            "請修改 VIDEO_SOURCE 為有效的影片路徑。"
        )

    # 確認模型存在
    if not Path(MODEL_PATH).exists():
        raise FileNotFoundError(
            f"找不到 SAM 3 模型：{MODEL_PATH}\n"
            "請至 https://huggingface.co/facebook/sam3 下載 sam3.pt 並放置於當前目錄。"
        )

    # 初始化 SAM3VideoSemanticPredictor
    overrides = dict(
        conf=CONFIDENCE,
        task="segment",
        mode="predict",
        imgsz=IMG_SIZE,
        model=MODEL_PATH,
        half=True,   # 使用 FP16 加速推理（GPU 需支援）
        save=False,  # 我們手動處理輸出，不使用自動儲存
    )
    predictor = SAM3VideoSemanticPredictor(overrides=overrides)

    # 取得影片基本資訊（用於設定輸出）
    cap = cv2.VideoCapture(VIDEO_SOURCE)
    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))
    cap.release()

    print(f"影片資訊：{width}x{height} @ {fps:.1f} FPS，共 {total_frames} 幀")

    # 建立 VideoWriter 輸出 output.mp4
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    writer = cv2.VideoWriter(OUTPUT_PATH, fourcc, fps, (width, height))

    # 執行文本提示視頻跟踪（stream=True 逐幀處理）
    print(f"\n開始分割追蹤...")
    results = predictor(
        source=VIDEO_SOURCE,
        text=TEXT_PROMPTS,
        stream=True,
    )

    frame_count = 0
    for r in results:
        frame_count += 1

        # 取得帶有分割遮罩的視覺化幀
        annotated_frame = r.plot()  # BGR numpy array

        # 寫入輸出影片
        writer.write(annotated_frame)

        # 顯示進度
        if frame_count % 10 == 0 or frame_count == 1:
            print(f"  處理第 {frame_count}/{total_frames} 幀 | 偵測到 {len(r.boxes) if r.boxes is not None else 0} 個物件")

    writer.release()

    print()
    print("=" * 60)
    print(f"✅ 完成！共處理 {frame_count} 幀")
    print(f"💾 輸出已儲存至：{OUTPUT_PATH}")
    print("=" * 60)


if __name__ == "__main__":
    main()