"""Vision-describe corpus images via OpenRouter."""

import base64
import io
import json
import os
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from threading import Lock

import requests
from PIL import Image

ENV_PATH = "/home/videobot/dev/.env"

DEFAULT_VISION_PROMPT = """Describe this image in detail. Include:
- WHO: characters depicted, their roles, recognizable figures, appearance/outfit
- WHAT: actions, gestures, interactions
- WHERE: setting, location, background elements
- TOPIC: theme or subject referenced
- STYLE: visual metaphors, symbolism, artistic style details
- TEXT: any text, labels, captions visible (transcribe in original language)
- COMPOSITION: layout, color palette, perspective, medium
If the page is text-only with no illustrations, respond with just: TEXT_ONLY"""

DEFAULT_MODEL = "anthropic/claude-haiku-4.5"
MAX_IMAGE_WIDTH = 1500
DEFAULT_WORKERS = 10


def _load_api_key():
    key = os.environ.get("OPENROUTER_API_KEY")
    if key:
        return key
    with open(ENV_PATH, "r") as f:
        for line in f:
            if line.startswith("OPENROUTER_API_KEY="):
                return line.strip().split("=", 1)[1]
    raise RuntimeError("OPENROUTER_API_KEY not found")


def _encode_image(image_path: str, max_width: int = None) -> tuple:
    """Load image, optionally resize, return (base64_str, media_type)."""
    from PIL import ImageFile
    ImageFile.LOAD_TRUNCATED_IMAGES = True
    img = Image.open(image_path)
    img.load()

    if max_width and img.width > max_width:
        ratio = max_width / img.width
        new_size = (max_width, int(img.height * ratio))
        img = img.resize(new_size, Image.LANCZOS)

    if img.mode in ("RGBA", "P", "LA"):
        img = img.convert("RGB")

    buf = io.BytesIO()
    img.save(buf, format="JPEG", quality=85)
    b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
    return b64, "image/jpeg"


def _describe_one_image(api_key: str, image_path: str, source: str,
                        extra_context: str = "",
                        model: str = DEFAULT_MODEL,
                        vision_prompt: str = DEFAULT_VISION_PROMPT):
    """Call vision model to describe a single image."""
    max_w = MAX_IMAGE_WIDTH if source == "djvu" else None
    b64, media_type = _encode_image(image_path, max_width=max_w)

    prompt = vision_prompt
    if extra_context:
        prompt += f"\n\nAdditional context about this image: {extra_context}"

    payload = {
        "model": model,
        "max_tokens": 1024,
        "messages": [{
            "role": "user",
            "content": [
                {"type": "image_url",
                 "image_url": {"url": f"data:{media_type};base64,{b64}"}},
                {"type": "text", "text": prompt},
            ],
        }],
    }

    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }

    for attempt in range(3):
        try:
            resp = requests.post(
                "https://openrouter.ai/api/v1/chat/completions",
                headers=headers, json=payload, timeout=90,
            )
            if resp.status_code == 429:
                time.sleep(2 ** (attempt + 1))
                continue
            if resp.status_code != 200:
                if attempt < 2:
                    time.sleep(2)
                    continue
                return None
            data = resp.json()
            content = data["choices"][0]["message"]["content"]
            return content.strip()
        except (requests.RequestException, KeyError, json.JSONDecodeError):
            if attempt < 2:
                time.sleep(2)
            else:
                return None
    return None


def _process_one(api_key: str, entry: dict, images_base: str,
                 model: str, vision_prompt: str):
    """Process a single catalog entry. Returns (id, description, is_text_only, error)."""
    image_path = str(Path(images_base) / entry["image_path"])
    if not os.path.exists(image_path):
        return (entry["id"], None, False, "MISSING")

    context_parts = []
    if entry.get("title"):
        context_parts.append(f"Title: {entry['title']}")
    if entry.get("year"):
        context_parts.append(f"Year: {entry['year']}")
    if entry.get("materials"):
        context_parts.append(f"Materials: {entry['materials']}")
    if entry.get("index_commentary"):
        context_parts.append(f"Commentary: {entry['index_commentary'][:300]}")
    extra = "; ".join(context_parts) if context_parts else ""

    try:
        description = _describe_one_image(
            api_key, image_path, entry.get("source", ""),
            extra, model, vision_prompt)
    except Exception as e:
        return (entry["id"], None, False, str(e))

    if description is None:
        return (entry["id"], None, False, "API_ERROR")
    if description.strip() == "TEXT_ONLY":
        return (entry["id"], None, True, None)

    return (entry["id"], description, False, None)


def describe_images(catalog_path: str, images_base: str,
                    model: str = DEFAULT_MODEL,
                    workers: int = DEFAULT_WORKERS,
                    limit: int = 0,
                    source_filter: str = None,
                    force: bool = False,
                    vision_prompt: str = None,
                    api_key: str = None):
    """Vision-describe all images in catalog that lack descriptions.

    Args:
        catalog_path: Path to catalog.json
        images_base: Base directory for resolving image_path in catalog entries
        model: Vision model to use
        workers: Number of parallel API workers
        limit: Max images to process (0=all)
        source_filter: Only process entries with this source value
        force: Re-describe even if description exists
        vision_prompt: Custom vision prompt (uses default if None)
        api_key: OpenRouter API key (auto-loaded if None)
    """
    if api_key is None:
        api_key = _load_api_key()
    if vision_prompt is None:
        vision_prompt = DEFAULT_VISION_PROMPT

    with open(catalog_path, "r", encoding="utf-8") as f:
        catalog = json.load(f)

    to_process = []
    for entry in catalog:
        if not force:
            if entry.get("description") or entry.get("skip"):
                continue
        if source_filter and entry.get("source") != source_filter:
            continue
        to_process.append(entry)

    if limit:
        to_process = to_process[:limit]

    already_done = sum(1 for e in catalog if e.get("description") or e.get("skip"))
    print(f"Catalog: {len(catalog)} total, {already_done} already done, "
          f"{len(to_process)} to process")
    print(f"Model: {model}, Workers: {workers}")

    if not to_process:
        print("Nothing to process.")
        return

    id_to_idx = {e["id"]: i for i, e in enumerate(catalog)}
    processed = 0
    errors = 0
    text_only = 0
    save_every = 50
    catalog_lock = Lock()
    start_time = time.time()

    def save_catalog():
        with open(catalog_path, "w", encoding="utf-8") as f:
            json.dump(catalog, f, ensure_ascii=False, indent=2)

    with ThreadPoolExecutor(max_workers=workers) as executor:
        futures = {}
        for entry in to_process:
            fut = executor.submit(
                _process_one, api_key, entry, images_base, model, vision_prompt)
            futures[fut] = entry

        for fut in as_completed(futures):
            entry = futures[fut]
            try:
                entry_id, description, is_text_only, error = fut.result()
            except Exception as e:
                errors += 1
                print(f"  EXCEPTION {entry['id']}: {e}")
                continue

            with catalog_lock:
                idx = id_to_idx[entry_id]
                if error:
                    errors += 1
                    print(f"  {error}: {entry_id}")
                elif is_text_only:
                    catalog[idx]["skip"] = True
                    catalog[idx]["description"] = None
                    text_only += 1
                else:
                    catalog[idx]["description"] = description
                    processed += 1

                total_done = processed + text_only + errors
                elapsed = time.time() - start_time
                rate = total_done / elapsed if elapsed > 0 else 0
                eta = (len(to_process) - total_done) / rate if rate > 0 else 0

                if total_done % 10 == 0 or total_done == len(to_process):
                    print(f"  [{total_done}/{len(to_process)}] "
                          f"{processed} ok, {text_only} text, {errors} err | "
                          f"{rate:.1f}/s, ETA {eta/60:.0f}m")

                if (processed + text_only) % save_every == 0:
                    save_catalog()

    save_catalog()
    elapsed = time.time() - start_time

    print(f"\nDone in {elapsed/60:.1f}m: {processed} described, "
          f"{text_only} text-only, {errors} errors")
    print(f"Rate: {(processed + text_only) / elapsed:.1f} images/sec")
    print(f"Total with descriptions: "
          f"{sum(1 for e in catalog if e.get('description'))}")
