#!/usr/bin/env python
"""Check screenshot quality via OpenAI gpt-5.4-mini vision.

Uses OPENAI_API_KEY from environment.
Analyzes a screenshot and returns structured JSON answering:
- Is the page fully rendered (not blank/loading/error)?
- Is real content visible (text, images, products)?
- Does the data look useful for an AI crawler's output?
- What page title/heading is visible?
"""

from __future__ import annotations

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

from openai import OpenAI

MODEL = "gpt-5.4-mini"


def encode_image(path: Path) -> str:
    return base64.b64encode(path.read_bytes()).decode("utf-8")


def analyze_screenshot(image_path: str | Path) -> dict:
    image_path = Path(image_path)
    if not image_path.exists():
        return {"error": f"File not found: {image_path}", "path": str(image_path)}

    api_key = os.environ.get("OPENAI_API_KEY", "")
    if not api_key:
        return {"error": "OPENAI_API_KEY not set in environment", "path": str(image_path)}

    client = OpenAI(api_key=api_key)
    image_b64 = encode_image(image_path)

    prompt = """Analyze this screenshot of a crawled webpage. Answer:

1. is_fully_rendered: Is the page fully loaded and rendered (not blank, white, error page, or stuck on a loading spinner)?
2. has_visible_content: Does the page show actual content like text, headings, product images, navigation menus? Not just a cookie banner or empty shell.
3. is_useful_crawl_output: Based on what's visible, would an AI crawler's output from this page be useful? Consider: readable text, clear structure, product information, navigation elements.
4. page_title: What is the main title or heading visible on the page?
5. has_images: Are there visible product images, photos, or graphical elements on the page?
6. content_quality: Rate the content quality 1-5 (1=empty/garbage, 5=rich well-structured content)
7. summary: One sentence describing what the page shows.

Return ONLY valid JSON, no markdown:"""

    started = time.perf_counter()
    try:
        response = client.chat.completions.create(
            model=MODEL,
            messages=[
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": prompt},
                        {
                            "type": "image_url",
                            "image_url": {"url": f"data:image/png;base64,{image_b64}", "detail": "low"},
                        },
                    ],
                }
            ],
            max_completion_tokens=500,
            temperature=0.0,
        )
        elapsed = time.perf_counter() - started
        text = response.choices[0].message.content
        result = json.loads(text)
        result["_elapsed_s"] = round(elapsed, 2)
        result["_model"] = MODEL
        result["_path"] = str(image_path)
        return result
    except Exception as exc:
        return {"error": f"OpenAI API error: {exc}", "path": str(image_path), "_elapsed_s": round(time.perf_counter() - started, 2)}


def check_all_screenshots(screenshots_dir: str | Path) -> list[dict]:
    screenshots_dir = Path(screenshots_dir)
    if not screenshots_dir.exists():
        return [{"error": f"Directory not found: {screenshots_dir}"}]

    pngs = sorted(screenshots_dir.glob("*tile*.png"))
    results = []
    for png in pngs:
        print(f"Checking {png.name}...")
        result = analyze_screenshot(png)
        results.append(result)
    return results


def main() -> int:
    import argparse
    p = argparse.ArgumentParser(description="Check screenshot quality via OpenAI vision")
    p.add_argument("image", nargs="?", help="Path to a single screenshot PNG")
    p.add_argument("--dir", help="Check all tile screenshots in directory")
    p.add_argument("--json", action="store_true", help="Output raw JSON only")
    args = p.parse_args()

    if args.dir:
        results = check_all_screenshots(args.dir)
        if args.json:
            print(json.dumps(results, ensure_ascii=False, indent=2))
            return 0

        good = sum(1 for r in results if r.get("is_useful_crawl_output"))
        bad = len(results) - good
        print(f"\n{len(results)} screenshots checked: {good} useful, {bad} not useful")
        for i, r in enumerate(results):
            status = "PASS" if r.get("is_useful_crawl_output") else "FAIL"
            title = r.get("page_title", "?") or "?"
            quality = r.get("content_quality", "?")
            print(f"  [{status}] {r.get('_path', f'#{i}')}: q={quality}/5 title={title[:60]}")
        return 0 if bad == 0 else 1

    if not args.image:
        p.print_help()
        return 1

    result = analyze_screenshot(args.image)
    if args.json:
        print(json.dumps(result, ensure_ascii=False, indent=2))
        return 0

    if "error" in result:
        print(f"ERROR: {result['error']}")
        return 1

    print(f"Fully rendered:   {'YES' if result.get('is_fully_rendered') else 'NO'}")
    print(f"Visible content:  {'YES' if result.get('has_visible_content') else 'NO'}")
    print(f"Useful output:    {'YES' if result.get('is_useful_crawl_output') else 'NO'}")
    print(f"Has images:       {'YES' if result.get('has_images') else 'NO'}")
    print(f"Quality:          {result.get('content_quality', '?')}/5")
    print(f"Title:            {result.get('page_title', '(none)')}")
    print(f"Summary:          {result.get('summary', '(none)')}")
    print(f"Elapsed:          {result.get('_elapsed_s')}s")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
