#!/opt/homebrew/bin/python3.11
"""
Interval VLM Caption - Analyze video at regular intervals (10s default)

Extracts frames on-the-fly with ffmpeg (no pre-storage) and analyzes with VLM.
Runs in background, non-blocking.

Usage:
    python interval_vlm_caption.py --file-uuid abc123 --video /path/to/video.mp4
    python interval_vlm_caption.py --file-uuid abc123 --video /path/to/video.mp4 --interval 10

Output:
    {output_dir}/{uuid}_interval_profile.json
"""

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

try:
    import requests
except ImportError:
    print("requests not installed: pip install requests", file=sys.stderr)
    sys.exit(1)


def get_video_duration(video_path: str) -> float:
    """Get video duration in seconds using ffprobe."""
    cmd = [
        "ffprobe", "-v", "quiet",
        "-show_entries", "format=duration",
        "-of", "json",
        video_path
    ]
    
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        return 0.0
    
    data = json.loads(result.stdout)
    return float(data["format"]["duration"])


def extract_frame_at_time(video_path: str, timestamp_sec: float, output_path: str) -> bool:
    """Extract a single frame at specific timestamp."""
    cmd = [
        "ffmpeg", "-y", "-v", "quiet",
        "-ss", str(timestamp_sec),
        "-i", video_path,
        "-vframes", "1",
        "-q:v", "2",
        output_path
    ]
    
    result = subprocess.run(cmd, capture_output=True)
    return result.returncode == 0


def call_vlm(image_path: str, prompt: str, model: str = "llava:7b", ollama_url: str = "http://localhost:11434") -> str:
    """Call Ollama VLM API."""
    with open(image_path, "rb") as f:
        image_b64 = base64.b64encode(f.read()).decode("utf-8")
    
    payload = {
        "model": model,
        "prompt": prompt,
        "images": [image_b64],
        "stream": False,
        "options": {"num_predict": 100}
    }
    
    try:
        resp = requests.post(f"{ollama_url}/api/generate", json=payload, timeout=30)
        resp.raise_for_status()
        data = resp.json()
        return data.get("response", "").strip()
    except Exception as e:
        print(f"[vlm] API error: {e}", file=sys.stderr)
        return ""


def get_embedding(text: str, model: str = "nomic-embed-text-v2-moe", ollama_url: str = "http://localhost:11434") -> list:
    """Get embedding from Ollama."""
    try:
        resp = requests.post(
            f"{ollama_url}/api/embed",
            json={"model": model, "input": text},
            timeout=30,
        )
        resp.raise_for_status()
        data = resp.json()
        return data.get("embeddings", [[]])[0]
    except Exception as e:
        print(f"[vlm] Embedding error: {e}", file=sys.stderr)
        return []


def store_to_qdrant(analysis: dict, file_uuid: str, interval_index: int, timestamp_sec: float, qdrant_url: str = "http://localhost:6333", qdrant_api_key: str = None) -> bool:
    """Store VLM results to Qdrant _vlm collection."""
    description = analysis.get("vlm_description", "")
    if not description:
        return False
    
    # Get embedding
    embedding = get_embedding(description)
    if not embedding:
        print(f"[vlm] Failed to get embedding for interval_{interval_index}", file=sys.stderr)
        return False
    
    # Generate point ID
    import hashlib
    point_id = int(hashlib.md5(f"{file_uuid}_interval_{interval_index}".encode()).hexdigest()[:16], 16)
    
    # Build payload
    payload = {
        "type": "interval",
        "file_uuid": file_uuid,
        "interval_index": interval_index,
        "timestamp_sec": timestamp_sec,
        **analysis,
    }
    
    # Upsert to Qdrant
    try:
        headers = {}
        if qdrant_api_key:
            headers["api-key"] = qdrant_api_key
        
        resp = requests.put(
            f"{qdrant_url}/collections/_vlm/points?wait=true",
            json={
                "points": [{
                    "id": point_id,
                    "vector": embedding,
                    "payload": payload,
                }]
            },
            headers=headers,
            timeout=30,
        )
        resp.raise_for_status()
        print(f"[vlm] Stored to Qdrant: interval_{interval_index}")
        return True
    except Exception as e:
        print(f"[vlm] Qdrant error: {e}", file=sys.stderr)
        return False


def analyze_frame(image_path: str, model: str = "llava:7b") -> dict:
    """Analyze a single frame with VLM - scene/background focus."""
    # Prompt 1: Scene description (concise)
    desc_prompt = "Describe this scene in one sentence. Focus on: location, main activity, visible objects. If unclear, say 'unclear'. Do not guess."
    description = call_vlm(image_path, desc_prompt, model)
    
    # Prompt 2: Classification (JSON)
    class_prompt = "Classify the scene. Answer in JSON: {\"location\": \"indoor/outdoor/unknown\", \"setting\": \"office/street/home/nature/studio/unknown\", \"lighting\": \"day/night/indoor-light/mixed/unknown\"}. Use 'unknown' if uncertain."
    class_raw = call_vlm(image_path, class_prompt, model)
    
    class_data = {}
    try:
        class_clean = class_raw.replace("```json", "").replace("```", "").strip()
        class_data = json.loads(class_clean)
    except:
        class_data = {}
    
    # Prompt 3: People count
    people_prompt = "How many people? Answer a number or 'unclear'."
    people_count = call_vlm(image_path, people_prompt, model).strip()
    
    # Prompt 4: Tags
    tags_prompt = "List 3 tags for this scene, comma-separated. Examples: office, street, crowd, nature."
    tags_raw = call_vlm(image_path, tags_prompt, model)
    tags = [t.strip() for t in tags_raw.replace(",", " ").split() if t.strip()][:3]
    
    return {
        "vlm_description": description,
        "vlm_location": class_data.get("location", "unknown"),
        "vlm_setting": class_data.get("setting", "unknown"),
        "vlm_lighting": class_data.get("lighting", "unknown"),
        "vlm_people_count": people_count,
        "vlm_tags": tags,
    }


def analyze_intervals(
    file_uuid: str,
    video_path: str,
    interval_sec: float = 10.0,
    output_dir: str = "/Users/accusys/momentry/output",
    model: str = "llava:7b",
    store_qdrant: bool = True,
) -> dict:
    """
    Analyze video at regular intervals.
    
    Returns:
        Summary dict with all interval analyses
    """
    # Get video duration
    duration = get_video_duration(video_path)
    if duration <= 0:
        print(f"[vlm] Cannot get video duration: {video_path}", file=sys.stderr)
        return {"error": "Cannot get duration"}
    
    print(f"[vlm] Video: {duration:.1f}s, interval: {interval_sec}s")
    
    # Calculate timestamps
    timestamps = []
    t = 0.0
    while t < duration:
        timestamps.append(t)
        t += interval_sec
    
    print(f"[vlm] Total frames to analyze: {len(timestamps)}")
    
    results = []
    qdrant_api_key = os.environ.get("QDRANT_API_KEY")
    
    # Create temp directory for frames
    with tempfile.TemporaryDirectory() as tmpdir:
        for i, ts in enumerate(timestamps):
            frame_path = f"{tmpdir}/frame_{i:04d}.jpg"
            
            # Extract frame
            success = extract_frame_at_time(video_path, ts, frame_path)
            if not success:
                print(f"[vlm] Failed to extract frame at {ts:.1f}s", file=sys.stderr)
                continue
            
            # Analyze
            print(f"[vlm] [{i+1}/{len(timestamps)}] {ts:.1f}s...", end=" ", flush=True)
            start_time = time.time()
            
            analysis = analyze_frame(frame_path, model)
            
            # Store to Qdrant
            if store_qdrant:
                store_to_qdrant(analysis, file_uuid, i, ts, qdrant_api_key=qdrant_api_key)
            
            elapsed = time.time() - start_time
            print(f"done ({elapsed:.1f}s)")
            
            results.append({
                "interval_index": i,
                "timestamp_sec": round(ts, 1),
                **analysis,
            })
    
    # Save profile
    profile = {
        "file_uuid": file_uuid,
        "video_duration_sec": round(duration, 1),
        "interval_sec": interval_sec,
        "total_intervals": len(timestamps),
        "analyzed": len(results),
        "model": model,
        "intervals": results,
    }
    
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)
    
    profile_path = output_path / f"{file_uuid}_interval_profile.json"
    with open(profile_path, "w") as f:
        json.dump(profile, f, indent=2)
    
    print(f"[vlm] Saved: {profile_path}")
    
    return profile


def main():
    parser = argparse.ArgumentParser(description="VLM analysis at regular intervals")
    parser.add_argument("--file-uuid", "-u", required=True, help="File UUID")
    parser.add_argument("--video", "-v", required=True, help="Video file path")
    parser.add_argument("--interval", "-i", type=float, default=10.0, help="Interval in seconds (default: 10)")
    parser.add_argument("--output-dir", "-o", default="/Users/accusys/momentry/output", help="Output directory")
    parser.add_argument("--model", "-m", default="llava:7b", help="VLM model name")
    parser.add_argument("--json", "-j", action="store_true", help="Output as JSON")
    args = parser.parse_args()
    
    result = analyze_intervals(
        args.file_uuid,
        args.video,
        args.interval,
        args.output_dir,
        args.model,
    )
    
    if args.json:
        print(json.dumps(result, indent=2))
    else:
        print(f"Analyzed {result.get('analyzed', 0)} intervals")


if __name__ == "__main__":
    main()