#!/usr/bin/env python3
"""Batch visual QA of rendered document pages/slides via direct vLLM API calls.

Bypasses the vision_analyze tool — useful when:
- vision_analyze fails due to config caching / model resolution issues
- You need to QA many images efficiently in one pass
- The active chat model is text-only but a multimodal model is on the same endpoint

Usage:
    python batch_vision_qa.py /path/to/images/ <num_images> [options]

Environment variables:
    NORIS_API_KEY   - API key for the vLLM endpoint (required)
    VISION_BASE_URL - Full chat completions URL (default: https://ai.noris.de/v1/chat/completions)
    VISION_MODEL    - Model name (default: vllm/release/gemma-4-31b-it)

Example:
    NORIS_API_KEY=sk-bf-xxx python batch_vision_qa.py /home/debian/ai101/qa/ 17

Expects files named slide-01.jpg, slide-02.jpg, ..., slide-NN.jpg.
Outputs per-slide status and a summary of any issues found.

Before using: verify the model accepts images (see references/model-capabilities.md).
"""

import sys
import os
import base64
import json
import time
import urllib.request

API_KEY = os.environ.get("NORIS_API_KEY", "")
BASE_URL = os.environ.get(
    "VISION_BASE_URL", "https://ai.noris.de/v1/chat/completions"
)
MODEL = os.environ.get("VISION_MODEL", "vllm/release/gemma-4-31b-it")

# Expanded 6-category prompt: overlap, overflow, contrast, layout, content, symbols
QA_PROMPT_TEMPLATE = """You are a visual QA reviewer for a professional presentation slide.
Analyze this slide carefully for ALL of these issues:
1. TEXT OVERLAP: Any text overlapping other text or graphics?
2. TEXT OVERFLOW: Any text extending beyond edges or getting cut off?
3. CONTRAST: Any low contrast between text and background making text hard to read?
4. LAYOUT: Misalignment, elements too close, insufficient margins, awkward spacing?
5. CONTENT: Unclear messaging, missing information, placeholder text, spelling errors, logical inconsistencies?
6. SYMBOLS: Any broken/missing Unicode characters showing as boxes or empty circles?

Respond in this exact JSON format:
{{"slide": {slide_num}, "overlap": "none|<description>", "overflow": "none|<description>", "contrast": "ok|<description>", "layout": "ok|<description>", "content": "ok|<description>", "symbols": "ok|<description>", "overall": "clean|issues found"}}

Be precise and thorough. If no issues, say "clean". If issues exist, describe them specifically."""


def qa_single_image(image_path: str, slide_num: int) -> dict:
    """Send one image to the vision model and return parsed QA result."""
    with open(image_path, "rb") as f:
        img_b64 = base64.b64encode(f.read()).decode()

    payload = {
        "model": MODEL,
        "messages": [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": QA_PROMPT_TEMPLATE.format(slide_num=slide_num),
                    },
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:image/jpeg;base64,{img_b64}"
                        },
                    },
                ],
            }
        ],
        "max_tokens": 500,
    }

    req = urllib.request.Request(
        BASE_URL,
        data=json.dumps(payload).encode(),
        headers={
            "Authorization": f"Bearer {API_KEY}",
            "Content-Type": "application/json",
        },
    )

    resp = urllib.request.urlopen(req, timeout=60)
    data = json.loads(resp.read())
    content = data["choices"][0]["message"].get("content", "")

    # Strip markdown code fences if present
    content = content.strip()
    if content.startswith("```"):
        content = content.split("\n", 1)[-1].rsplit("```", 1)[0].strip()

    try:
        return json.loads(content)
    except json.JSONDecodeError:
        return {
            "slide": slide_num,
            "overlap": "parse error",
            "overflow": "parse error",
            "contrast": "parse error",
            "layout": "parse error",
            "content": "parse error",
            "symbols": "parse error",
            "overall": content[:200],
        }


def main():
    if len(sys.argv) < 3:
        print(f"Usage: {sys.argv[0]} <images_dir> <num_images>")
        print(f"Example: {sys.argv[0]} /home/debian/ai101/qa/ 17")
        sys.exit(1)

    images_dir = sys.argv[1]
    num_slides = int(sys.argv[2])

    if not API_KEY:
        print("ERROR: Set NORIS_API_KEY environment variable.")
        sys.exit(1)

    results = []
    issues_found = []

    for i in range(1, num_slides + 1):
        slide_path = os.path.join(images_dir, f"slide-{i:02d}.jpg")
        if not os.path.exists(slide_path):
            print(f"Slide {i}: file not found — {slide_path}")
            continue

        try:
            result = qa_single_image(slide_path, i)
            results.append(result)
            status = result.get("overall", "?")
            marker = "⚠️" if "issue" in status.lower() else "✅"
            print(f"Slide {i}/{num_slides} {marker} {status}")
            if "issue" in status.lower():
                issues_found.append(result)
        except Exception as e:
            print(f"Slide {i} FAILED: {e}")
            results.append({"slide": i, "error": str(e)})

        time.sleep(0.5)

    print("\n" + "=" * 60)
    if issues_found:
        print(f"ISSUES ON {len(issues_found)} SLIDE(S):")
        for r in issues_found:
            print(f"  Slide {r['slide']}:")
            for key in ["overlap", "overflow", "contrast", "layout", "content", "symbols"]:
                val = r.get(key, "")
                if val and val not in ("none", "ok", ""):
                    print(f"    {key}: {val}")
    else:
        print("ALL SLIDES CLEAN ✅")


if __name__ == "__main__":
    main()