# -*- coding: utf-8 -*-
"""
run.py - Entry point for SAM 2 Robust Tracker
==============================================
Usage examples:
  python run.py --video my_video.mp4 --prompt "person . car . dog"
  python run.py --video my_video.mp4 --prompt "person . car" --fps 5 --chunk 100
  python run.py --video my_video.mp4 --prompt prompts.txt
"""

import os
import sys
import argparse
import shutil
import torch

# Resolve SAM 2 checkpoint path
HOME_DIR          = os.path.expanduser("~")
DEFAULT_SAM2_CKPT = os.path.join(
    HOME_DIR, ".cache", "torch", "hub", "checkpoints", "sam2.1_hiera_small.pt"
)
DEFAULT_SAM2_CFG  = "configs/sam2.1/sam2.1_hiera_s.yaml"

# Default detection prompt
DEFAULT_PROMPT = (
    "# Structures & Openings\n"
    "room door, door handle, wardrobe, window, stairs, socket, light switch\n\n"
    "# Furniture\n"
    "sofa, armchair, chair, table, bed, cabinet, bookcase, shelf, desk, nightstand\n\n"
    "# Kitchen & Bathroom\n"
    "sink, faucet, refrigerator, stove, oven, microwave, toilet, bathtub, shower\n\n"
    "# Lighting & Decor\n"
    "curtain, rug, lamp, painting, vase, clock, plant\n\n"
    "# Electronics\n"
    "tv, monitor, air conditioner, fan, laptop, phone, remote control"
)


# ──────────────────────────────────────────────────────────────────────────────
def parse_args():
    p = argparse.ArgumentParser(
        description="SAM 2 + Grounding DINO Robust Object Tracker",
        formatter_class=argparse.RawTextHelpFormatter
    )
    p.add_argument("--video",    required=True,
                   help="Path to input video file.")
    p.add_argument("--prompt",   default=DEFAULT_PROMPT,
                   help="Detection prompt: comma/dot separated labels, or path to .txt file.")
    p.add_argument("--output",   default="tracked_output.mp4",
                   help="Path for output annotated video. (default: tracked_output.mp4)")
    p.add_argument("--fps",      type=float, default=2.0,
                   help="Target processing FPS. 2.0 = 1 frame every 0.5s. (default: 2.0)")
    p.add_argument("--max-size", type=int,   default=720,
                   help="Longest edge in pixels after resize. (default: 720)")
    p.add_argument("--chunk",    type=int,   default=120,
                   help="SAM 2 sliding window size in frames. (default: 120)\n"
                        "  4 GB VRAM → 80  |  6 GB → 120  |  8 GB+ → 200")
    p.add_argument("--box-thresh",  type=float, default=0.30,
                   help="DINO box confidence threshold. (default: 0.30)")
    p.add_argument("--text-thresh", type=float, default=0.25,
                   help="DINO text confidence threshold. (default: 0.25)")
    p.add_argument("--iou-thresh",  type=float, default=0.45,
                   help="NMS IoU threshold. (default: 0.45)")
    p.add_argument("--max-objects", type=int, default=15,
                   help="Max objects to track (VRAM guard). (default: 15)")
    p.add_argument("--stabilize",   action="store_true", default=False,
                   help="Apply ORB-based video stabilization before tracking.")
    p.add_argument("--blur-thresh", type=float, default=60.0,
                   help="Skip blurry frames below this Laplacian variance. 0 = off.")
    p.add_argument("--sam2-ckpt",   default=DEFAULT_SAM2_CKPT,
                   help=f"Path to SAM 2 checkpoint .pt file.\n(default: {DEFAULT_SAM2_CKPT})")
    p.add_argument("--sam2-cfg",    default=DEFAULT_SAM2_CFG,
                   help=f"SAM 2 model config yaml.\n(default: {DEFAULT_SAM2_CFG})")
    p.add_argument("--frames-dir",  default="__frames_temp__",
                   help="Temp directory for extracted frames. (default: __frames_temp__)")
    p.add_argument("--keep-frames", action="store_true", default=False,
                   help="Keep extracted frames on disk after tracking.")
    return p.parse_args()


# ──────────────────────────────────────────────────────────────────────────────
def load_prompt(prompt_arg: str) -> str:
    """If prompt_arg is a path to a .txt file, read it; otherwise return as-is."""
    if os.path.isfile(prompt_arg):
        with open(prompt_arg, "r", encoding="utf-8") as f:
            return f.read()
    return prompt_arg


# ──────────────────────────────────────────────────────────────────────────────
def print_banner(args):
    vram_str = (f"{torch.cuda.get_device_properties(0).total_memory // 1024**2} MB"
                if torch.cuda.is_available() else "CPU mode")
    print()
    print("=" * 65)
    print("   SAM 2 Robust Tracker  -  Grounding DINO + Sliding Window")
    print("=" * 65)
    print(f"  Input  : {args.video}")
    print(f"  Output : {args.output}")
    print(f"  FPS    : {args.fps}  (1 frame every {1/args.fps:.2f}s)")
    print(f"  Chunk  : {args.chunk} frames per SAM-2 window")
    print(f"  MaxObj : {args.max_objects}")
    print(f"  MaxPx  : {args.max_size}px")
    print(f"  Stab   : {args.stabilize}")
    print(f"  VRAM   : {vram_str}")
    print("=" * 65)
    print()


# ──────────────────────────────────────────────────────────────────────────────
def main():
    args   = parse_args()
    prompt = load_prompt(args.prompt)

    if not os.path.exists(args.video):
        print(f"[ERROR] Video not found: {args.video}")
        sys.exit(1)

    if not os.path.exists(args.sam2_ckpt):
        print(f"[ERROR] SAM 2 checkpoint not found: {args.sam2_ckpt}")
        print("  Download with: python download_models.py")
        sys.exit(1)

    print_banner(args)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"[Device] {device}\n")

    # ── Import our modules ───────────────────────────────────────────────────
    from tracker import (
        VideoFrameStore,
        DinoDetector,
        SAM2Tracker,
        TrackedObject,
    )

    # ── Step 1: Extract frames ───────────────────────────────────────────────
    print("-" * 50)
    print("Step 1/3 - Extract Frames")
    print("-" * 50)

    store = VideoFrameStore(
        video_path      = args.video,
        output_dir      = args.frames_dir,
        target_fps      = args.fps,
        max_size        = args.max_size,
        blur_threshold  = args.blur_thresh,
        stabilize       = args.stabilize,
    )
    n_frames = store.extract()

    if n_frames == 0:
        print("[ERROR] No frames extracted.")
        sys.exit(1)

    print(f"  → {n_frames} frames  Âx  {store.width}×{store.height}px\n")

    # ── Step 2: DINO detection on frame 0 ────────────────────────────────────
    print("-" * 50)
    print("Step 2/3 - Initial Object Detection (DINO)")
    print("-" * 50)

    dino = DinoDetector(device)
    dino.load()

    first_frame = store.frame_paths[0]
    boxes, scores, labels = dino.detect(
        first_frame, prompt,
        box_threshold  = args.box_thresh,
        text_threshold = args.text_thresh,
        iou_threshold  = args.iou_thresh,
    )

    if len(boxes) == 0:
        print("[WARN] DINO found no objects in the first frame.")
        print("  Tip: lower --box-thresh or widen your --prompt.")
        sys.exit(1)

    # Cap to max_objects (by score, highest first)
    if len(boxes) > args.max_objects:
        top_k  = scores.argsort()[::-1][:args.max_objects]
        boxes  = boxes[top_k]
        scores = scores[top_k]
        labels = [labels[i] for i in top_k]
        print(f"  [Cap] Keeping top {args.max_objects} objects by confidence.")

    # Build TrackedObject list
    tracked_objects = []
    for i, (box, score, label) in enumerate(zip(boxes, scores, labels)):
        obj = TrackedObject(obj_id=i, label=label, box=box)
        tracked_objects.append(obj)
        print(f"  [ID {i:2d}] {label:<30s} conf={score:.3f}  box={box.astype(int).tolist()}")

    print(f"\n  → {len(tracked_objects)} objects registered.\n")

    # ── Step 3: SAM 2 sliding-window tracking ────────────────────────────────
    print("-" * 50)
    print("Step 3/3 - SAM 2 Sliding-Window Tracking")
    print("-" * 50)

    sam2 = SAM2Tracker(
        sam2_checkpoint = args.sam2_ckpt,
        sam2_cfg        = args.sam2_cfg,
        device          = device,
        chunk_size      = args.chunk,
    )
    sam2.load()

    def progress(done, total):
        pct = 100 * done / total
        filled = int(pct // 2)
        bar = "#" * filled + "-" * (50 - filled)
        print(f"\r  [{bar}] {pct:5.1f}%  ({done}/{total} frames)", end="", flush=True)

    result_labels = sam2.track_video(
        frame_store     = store,
        tracked_objects = tracked_objects,
        dino            = dino,
        prompt          = prompt,
        box_threshold   = args.box_thresh,
        text_threshold  = args.text_thresh,
        iou_threshold   = args.iou_thresh,
        output_path     = args.output,
        progress_cb     = progress,
    )

    print()

    # ── Cleanup ──────────────────────────────────────────────────────────────
    if not args.keep_frames and os.path.exists(args.frames_dir):
        shutil.rmtree(args.frames_dir, ignore_errors=True)
        print(f"[Clean] Removed temp frames: {args.frames_dir}")

    # ── Summary ──────────────────────────────────────────────────────────────
    print()
    print("=" * 65)
    print("  DONE!")
    print("=" * 65)
    print(f"  Output video : {os.path.abspath(args.output)}")
    print(f"  Tracked      : {len(result_labels)} objects")
    for i, label in enumerate(result_labels):
        status = "OK" if not tracked_objects[i].lost else "LOST at end"
        print(f"    ID {i:2d} [{status}] {label}")
    print("=" * 65)


if __name__ == "__main__":
    main()