#!/usr/bin/env python3
"""
generate_scenes_from_media.py — Auto-generate scene descriptions from image/video assets.

Scans a directory for images/videos, generates scene descriptions using
a vision model, and outputs as training pairs in JSONL format.

Usage:
    python3 scripts/generate_scenes_from_media.py --assets ~/assets/ --output training-data/media-scenes.jsonl
    python3 scripts/generate_scenes_from_media.py --assets ~/assets/ --model llava --dry-run
    python3 scripts/generate_scenes_from_media.py --assets ~/assets/ --max 10 --json
"""

import argparse
import hashlib
import json
import os
import subprocess
import sys
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Dict, List, Optional, Tuple

# Supported media formats
IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp", ".tiff"}
VIDEO_EXTENSIONS = {".mp4", ".mov", ".avi", ".mkv", ".webm", ".flv"}
ALL_EXTENSIONS = IMAGE_EXTENSIONS | VIDEO_EXTENSIONS


def find_media_files(assets_dir: str, max_files: int = 0) -> List[Path]:
    """Scan directory for media files."""
    assets_path = Path(assets_dir)
    if not assets_path.exists():
        print(f"ERROR: Directory not found: {assets_dir}", file=sys.stderr)
        return []

    media_files = []
    for ext in sorted(ALL_EXTENSIONS):
        media_files.extend(assets_path.rglob(f"*{ext}"))
        media_files.extend(assets_path.rglob(f"*{ext.upper()}"))

    # Deduplicate
    media_files = sorted(set(media_files))

    if max_files > 0:
        media_files = media_files[:max_files]

    return media_files


def file_hash(filepath: Path) -> str:
    """Generate hash for file deduplication."""
    return hashlib.sha256(str(filepath).encode()).hexdigest()[:16]


def generate_description_prompt(filepath: Path) -> str:
    """Generate the prompt for vision model."""
    if filepath.suffix.lower() in IMAGE_EXTENSIONS:
        return (
            "Describe this image as a visual scene for a training dataset. "
            "Include: mood, dominant colors (2-3), composition type, camera angle, "
            "and a vivid 1-2 sentence description. Format as JSON with keys: "
            "mood, colors, composition, camera, description."
        )
    else:
        return (
            "Describe this video frame as a visual scene for a training dataset. "
            "Include: mood, dominant colors (2-3), composition type, camera movement, "
            "and a vivid 1-2 sentence description. Format as JSON with keys: "
            "mood, colors, composition, camera, description."
        )


def call_vision_model(filepath: Path, model: str = "llava") -> Optional[dict]:
    """
    Call a vision model to generate scene description.

    Supports:
    - llava (local via ollama)
    - gpt-4-vision (OpenAI API)
    - claude-vision (Anthropic API)
    """
    prompt = generate_description_prompt(filepath)

    try:
        if model.startswith("llava") or model == "ollama":
            # Local Ollama with LLaVA
            result = subprocess.run(
                ["curl", "-s", "http://localhost:11434/api/generate", "-d",
                 json.dumps({
                     "model": "llava",
                     "prompt": prompt,
                     "images": [str(filepath)],
                     "stream": False,
                 })],
                capture_output=True, text=True, timeout=60
            )
            if result.returncode == 0:
                response = json.loads(result.stdout)
                return parse_description(response.get("response", ""))

        elif model.startswith("gpt-4"):
            # OpenAI GPT-4 Vision (requires API key)
            import base64
            with open(filepath, "rb") as f:
                image_data = base64.b64encode(f.read()).decode()

            api_key = os.environ.get("OPENAI_API_KEY")
            if not api_key:
                print("ERROR: OPENAI_API_KEY not set", file=sys.stderr)
                return None

            result = subprocess.run(
                ["curl", "-s", "https://api.openai.com/v1/chat/completions",
                 "-H", f"Authorization: Bearer {api_key}",
                 "-H", "Content-Type: application/json",
                 "-d", json.dumps({
                     "model": "gpt-4-vision-preview",
                     "messages": [{
                         "role": "user",
                         "content": [
                             {"type": "text", "text": prompt},
                             {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image_data}"}}
                         ]
                     }],
                     "max_tokens": 500
                 })],
                capture_output=True, text=True, timeout=60
            )
            if result.returncode == 0:
                response = json.loads(result.stdout)
                content = response["choices"][0]["message"]["content"]
                return parse_description(content)

        elif model.startswith("claude"):
            # Anthropic Claude Vision (requires API key)
            import base64
            with open(filepath, "rb") as f:
                image_data = base64.b64encode(f.read()).decode()

            api_key = os.environ.get("ANTHROPIC_API_KEY")
            if not api_key:
                print("ERROR: ANTHROPIC_API_KEY not set", file=sys.stderr)
                return None

            media_type = "image/jpeg" if filepath.suffix.lower() in {".jpg", ".jpeg"} else "image/png"
            result = subprocess.run(
                ["curl", "-s", "https://api.anthropic.com/v1/messages",
                 "-H", f"x-api-key: {api_key}",
                 "-H", "anthropic-version: 2023-06-01",
                 "-H", "Content-Type: application/json",
                 "-d", json.dumps({
                     "model": "claude-3-opus-20240229",
                     "max_tokens": 500,
                     "messages": [{
                         "role": "user",
                         "content": [
                             {"type": "image", "source": {"type": "base64", "media_type": media_type, "data": image_data}},
                             {"type": "text", "text": prompt}
                         ]
                     }]
                 })],
                capture_output=True, text=True, timeout=60
            )
            if result.returncode == 0:
                response = json.loads(result.stdout)
                content = response["content"][0]["text"]
                return parse_description(content)

    except (subprocess.TimeoutExpired, json.JSONDecodeError, KeyError) as e:
        print(f"ERROR calling vision model: {e}", file=sys.stderr)

    return None


def parse_description(text: str) -> dict:
    """Parse model response into structured description."""
    # Try to extract JSON from response
    import re
    json_match = re.search(r'\{[^}]+\}', text, re.DOTALL)
    if json_match:
        try:
            return json.loads(json_match.group())
        except json.JSONDecodeError:
            pass

    # Fallback: parse manually
    desc = {
        "mood": "unknown",
        "colors": [],
        "composition": "unknown",
        "camera": "unknown",
        "description": text[:500],
    }

    # Try to extract mood
    mood_match = re.search(r'mood["\s:]+(\w+)', text, re.IGNORECASE)
    if mood_match:
        desc["mood"] = mood_match.group(1).lower()

    # Try to extract colors
    color_match = re.search(r'colors?["\s:]+\[([^\]]+)\]', text, re.IGNORECASE)
    if color_match:
        desc["colors"] = [c.strip().strip('"').strip("'") for c in color_match.group(1).split(",")]

    return desc


def generate_training_pair(filepath: Path, description: dict, model: str) -> dict:
    """Generate a training pair from media file and description."""
    return {
        "source_file": str(filepath),
        "source_hash": file_hash(filepath),
        "source_type": "media_asset",
        "media_type": "image" if filepath.suffix.lower() in IMAGE_EXTENSIONS else "video",
        "model": model,
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "source_session_id": f"media-gen-{int(time.time())}",
        "prompt": f"Describe the visual scene in {filepath.name}",
        "response": description.get("description", ""),
        "scene": {
            "mood": description.get("mood", "unknown"),
            "colors": description.get("colors", []),
            "composition": description.get("composition", "unknown"),
            "camera": description.get("camera", "unknown"),
            "description": description.get("description", ""),
        },
    }


def main():
    parser = argparse.ArgumentParser(description="Generate scene descriptions from media")
    parser.add_argument("--assets", required=True, help="Assets directory to scan")
    parser.add_argument("--output", help="Output JSONL file path")
    parser.add_argument("--model", default="llava", help="Vision model (llava/gpt-4/claude)")
    parser.add_argument("--max", type=int, default=0, help="Max files to process (0=all)")
    parser.add_argument("--dry-run", action="store_true", help="Don't call vision model")
    parser.add_argument("--json", action="store_true", help="JSON output")
    args = parser.parse_args()

    media_files = find_media_files(args.assets, args.max)
    if not media_files:
        print("No media files found.", file=sys.stderr)
        sys.exit(1)

    print(f"Found {len(media_files)} media files in {args.assets}")

    if args.dry_run:
        print("\nDry run — files to process:")
        for f in media_files[:20]:
            print(f"  {f.relative_to(args.assets)}")
        if len(media_files) > 20:
            print(f"  ... and {len(media_files) - 20} more")
        sys.exit(0)

    pairs = []
    errors = 0

    for i, filepath in enumerate(media_files, 1):
        print(f"[{i}/{len(media_files)}] Processing {filepath.name}...", end=" ", flush=True)

        description = call_vision_model(filepath, args.model)
        if description:
            pair = generate_training_pair(filepath, description, args.model)
            pairs.append(pair)
            print(f"OK (mood: {pair['scene']['mood']})")
        else:
            errors += 1
            print("ERROR")

    # Output
    output_path = args.output or "training-data/media-scene-descriptions.jsonl"
    if args.json:
        print(json.dumps({"pairs": pairs, "total": len(pairs), "errors": errors}, indent=2))
    else:
        with open(output_path, 'w') as f:
            for pair in pairs:
                f.write(json.dumps(pair, ensure_ascii=False) + '\n')

        print(f"\nGenerated {len(pairs)} scene descriptions ({errors} errors)")
        print(f"Output: {output_path}")


if __name__ == "__main__":
    main()