#!/usr/bin/env python3
"""
Benchmark ArcisVLM vs Qwen3-VL on real camera frames.

Captures frames from live cameras, runs the SAME questions through both models,
and produces a side-by-side comparison report.

Qwen3-VL-8B is loaded via transformers (HuggingFace).
ArcisVLM uses the local API.

Usage:
    # Run on GPU instance with both models
    python3 scripts/benchmark_vs_qwen.py \
        --arcisvlm-api http://localhost:8000 \
        --qwen-model Qwen/Qwen3-VL-8B-Instruct \
        --camera-frames data/test_frames/ \
        --output benchmark_results.json

    # Or capture frames from live cameras first
    python3 scripts/benchmark_vs_qwen.py \
        --arcisvlm-api http://localhost:8000 \
        --capture-from "rtsp://admin:@ATPL-900064-AIPTZ.torqueverse.dev:64/ch0_0.264" \
        --qwen-model Qwen/Qwen3-VL-8B-Instruct
"""

import argparse
import base64
import json
import os
import sys
import time
from pathlib import Path

# Test questions covering all 8 agent types
BENCHMARK_QUESTIONS = [
    # VQA
    {"question": "What is happening in this image?", "agent": "vqa", "category": "general"},
    {"question": "Is there a person in this image?", "agent": "vqa", "category": "yes_no"},
    {"question": "What color is the main vehicle?", "agent": "vqa", "category": "attribute"},

    # Detect
    {"question": "List all objects visible in this scene.", "agent": "detect", "category": "detection"},
    {"question": "What types of vehicles are present?", "agent": "detect", "category": "vehicle_detect"},

    # Caption
    {"question": "Describe this image in detail.", "agent": "caption", "category": "description"},
    {"question": "Write a detailed caption for this surveillance camera image.", "agent": "caption", "category": "surveillance_caption"},

    # Count
    {"question": "How many people are in this image?", "agent": "count", "category": "count_people"},
    {"question": "How many vehicles can you see?", "agent": "count", "category": "count_vehicles"},

    # OCR
    {"question": "What text or signs are visible in this image?", "agent": "ocr", "category": "text_reading"},

    # Alert
    {"question": "Is there any suspicious or unusual activity in this scene?", "agent": "alert", "category": "anomaly"},
    {"question": "Are there any security concerns visible?", "agent": "alert", "category": "security"},

    # Reason
    {"question": "Analyze this scene and describe what is happening, including any potential concerns.", "agent": "reason", "category": "analysis"},

    # Track (single frame — limited, but tests spatial understanding)
    {"question": "Describe the positions and movements of people in this scene.", "agent": "track", "category": "spatial"},
]


def capture_frames(camera_urls: list[str], output_dir: str, n_frames: int = 3):
    """Capture frames from live cameras."""
    import cv2
    os.makedirs(output_dir, exist_ok=True)
    frames = []

    for url in camera_urls:
        cam_id = url.split("/")[-1].split(".")[0][:20]
        print(f"  Capturing from {cam_id}...")
        cap = cv2.VideoCapture(url, cv2.CAP_FFMPEG)
        if not cap.isOpened():
            print(f"    FAILED to open {url}")
            continue

        for i in range(n_frames):
            ret, frame = cap.read()
            if ret:
                path = os.path.join(output_dir, f"{cam_id}_frame{i}.jpg")
                cv2.imwrite(path, frame)
                frames.append(path)
                print(f"    Saved: {path}")
            time.sleep(1)
        cap.release()

    return frames


def query_arcisvlm(api_url: str, image_path: str, question: str, task_type: str) -> dict:
    """Query ArcisVLM API with an image."""
    import requests

    with open(image_path, "rb") as f:
        img_b64 = base64.b64encode(f.read()).decode()

    try:
        resp = requests.post(
            f"{api_url}/api/v1/query",
            json={
                "question": question,
                "task_type": task_type,
                "image_base64": img_b64,
            },
            timeout=120,
        )
        if resp.ok:
            data = resp.json()
            return {
                "answer": data.get("answer", ""),
                "time_ms": data.get("processing_time_ms", 0),
                "expert": data.get("expert_used", ""),
                "output_type": data.get("output_type", "text"),
            }
    except Exception as e:
        return {"answer": f"ERROR: {e}", "time_ms": 0, "expert": "", "output_type": "error"}

    return {"answer": "API error", "time_ms": 0}


def query_qwen(model, processor, image_path: str, question: str, device: str = "cuda") -> dict:
    """Query Qwen3-VL with an image."""
    from PIL import Image
    import torch

    start = time.time()
    try:
        image = Image.open(image_path).convert("RGB")

        messages = [
            {"role": "user", "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": question},
            ]}
        ]

        text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        inputs = processor(text=[text], images=[image], return_tensors="pt", padding=True).to(device)

        with torch.no_grad():
            generated = model.generate(**inputs, max_new_tokens=256)

        output = processor.batch_decode(generated[:, inputs.input_ids.shape[1]:], skip_special_tokens=True)[0]
        elapsed = (time.time() - start) * 1000

        return {"answer": output.strip(), "time_ms": elapsed}
    except Exception as e:
        return {"answer": f"ERROR: {e}", "time_ms": (time.time() - start) * 1000}


def run_benchmark(arcisvlm_api: str, qwen_model_name: str, frame_paths: list[str],
                  device: str = "cuda", output_path: str = "benchmark_results.json"):
    """Run side-by-side benchmark."""

    print("=" * 80)
    print("  ArcisVLM vs Qwen3-VL Benchmark")
    print("=" * 80)

    # Load Qwen model
    print(f"\nLoading Qwen model: {qwen_model_name}...")
    from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
    import torch

    qwen_processor = AutoProcessor.from_pretrained(qwen_model_name)
    qwen_model = Qwen2VLForConditionalGeneration.from_pretrained(
        qwen_model_name,
        torch_dtype=torch.bfloat16,
        device_map="auto",
    )
    print(f"  Qwen loaded: {sum(p.numel() for p in qwen_model.parameters()):,} params")

    results = []

    for frame_idx, frame_path in enumerate(frame_paths[:5]):  # Max 5 frames
        print(f"\n{'='*60}")
        print(f"  Frame {frame_idx + 1}: {os.path.basename(frame_path)}")
        print(f"{'='*60}")

        for q in BENCHMARK_QUESTIONS:
            question = q["question"]
            agent = q["agent"]
            category = q["category"]

            print(f"\n  Q [{agent}]: {question[:60]}...")

            # ArcisVLM
            arc_result = query_arcisvlm(arcisvlm_api, frame_path, question, agent)
            print(f"    ArcisVLM: {arc_result['answer'][:100]}  [{arc_result['time_ms']:.0f}ms]")

            # Qwen
            qwen_result = query_qwen(qwen_model, qwen_processor, frame_path, question, device)
            print(f"    Qwen3-VL: {qwen_result['answer'][:100]}  [{qwen_result['time_ms']:.0f}ms]")

            results.append({
                "frame": os.path.basename(frame_path),
                "question": question,
                "agent": agent,
                "category": category,
                "arcisvlm_answer": arc_result["answer"],
                "arcisvlm_time_ms": arc_result["time_ms"],
                "qwen_answer": qwen_result["answer"],
                "qwen_time_ms": qwen_result["time_ms"],
            })

    # Summary
    print(f"\n{'='*80}")
    print("  SUMMARY")
    print(f"{'='*80}")

    arc_avg_time = sum(r["arcisvlm_time_ms"] for r in results) / max(len(results), 1)
    qwen_avg_time = sum(r["qwen_time_ms"] for r in results) / max(len(results), 1)

    print(f"  Total questions: {len(results)}")
    print(f"  ArcisVLM avg time: {arc_avg_time:.0f}ms")
    print(f"  Qwen3-VL avg time: {qwen_avg_time:.0f}ms")

    # Save
    with open(output_path, "w") as f:
        json.dump({
            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S UTC", time.gmtime()),
            "arcisvlm_api": arcisvlm_api,
            "qwen_model": qwen_model_name,
            "total_questions": len(results),
            "arcisvlm_avg_time_ms": arc_avg_time,
            "qwen_avg_time_ms": qwen_avg_time,
            "results": results,
        }, f, indent=2)
    print(f"\n  Results saved to: {output_path}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="ArcisVLM vs Qwen3-VL Benchmark")
    parser.add_argument("--arcisvlm-api", default="http://localhost:8000")
    parser.add_argument("--qwen-model", default="Qwen/Qwen3-VL-8B-Instruct")
    parser.add_argument("--camera-frames", default="data/test_frames",
                        help="Directory with test frame images")
    parser.add_argument("--capture-from", nargs="*", default=None,
                        help="Camera URLs to capture frames from")
    parser.add_argument("--device", default="cuda")
    parser.add_argument("--output", default="benchmark_results.json")
    args = parser.parse_args()

    # Capture frames if requested
    if args.capture_from:
        print("Capturing frames from cameras...")
        frames = capture_frames(args.capture_from, args.camera_frames)
    else:
        frames = sorted(Path(args.camera_frames).glob("*.jpg"))
        if not frames:
            print(f"No frames found in {args.camera_frames}. Use --capture-from to capture.")
            sys.exit(1)
        frames = [str(f) for f in frames]

    print(f"Using {len(frames)} test frames")
    run_benchmark(args.arcisvlm_api, args.qwen_model, frames, args.device, args.output)