#!/usr/bin/env python3
"""aria-vision-helper — Gemini Vision für Image-Analyse (Phase 2 Multi-Modal-Helper).

Use-Cases:
- Telegram-Attachments mit Image: schnelle Beschreibung
- OCR auf Screenshots
- Chart-Analyse, Architektur-Diagramme
- claude-watch Frame-Synthese (statt Claude-Multi-Modal $)

Gemini 2.5-flash ist günstig + schnell für Image-Inputs. Pro für komplexe Analysen.

CLI:
    python3 aria-vision-helper.py /path/to/image.jpg
    python3 aria-vision-helper.py image.jpg --task "OCR all visible text"
    python3 aria-vision-helper.py screenshot.png --pro --task "explain the architecture diagram"
"""
from __future__ import annotations

import argparse
import base64
import json
import os
import sys
import time
import urllib.request
from pathlib import Path

sys.path.insert(0, "/root/aria/scripts")

# Load Gemini-Key
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "")
if not GEMINI_API_KEY:
    sec = Path("/root/.aria-secrets/gemini.env")
    if sec.exists():
        for line in sec.read_text().splitlines():
            if line.startswith("GEMINI_API_KEY="):
                GEMINI_API_KEY = line.split("=", 1)[1].strip().strip('"').strip("'")
                break

MIME_BY_EXT = {
    ".jpg": "image/jpeg", ".jpeg": "image/jpeg",
    ".png": "image/png", ".webp": "image/webp", ".gif": "image/gif",
}


def analyze_image(path: Path, *, task: str = "Describe this image in detail.",
                   model: str = "gemini-2.5-flash", max_tokens: int = 1500) -> dict:
    if not GEMINI_API_KEY:
        raise SystemExit("ERROR: GEMINI_API_KEY missing. /root/.aria-secrets/gemini.env")
    if not path.exists():
        raise SystemExit(f"ERROR: image not found: {path}")

    mime = MIME_BY_EXT.get(path.suffix.lower(), "image/jpeg")
    image_b64 = base64.b64encode(path.read_bytes()).decode("ascii")

    body = {
        "contents": [{
            "parts": [
                {"text": task},
                {"inline_data": {"mime_type": mime, "data": image_b64}},
            ],
        }],
        "generationConfig": {"maxOutputTokens": max_tokens},
    }
    url = f"https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent?key={GEMINI_API_KEY}"
    t0 = time.monotonic()
    req = urllib.request.Request(url, data=json.dumps(body).encode(),
                                  headers={"Content-Type": "application/json"}, method="POST")
    with urllib.request.urlopen(req, timeout=60) as resp:
        result = json.loads(resp.read())
    latency_ms = int((time.monotonic() - t0) * 1000)

    candidate = (result.get("candidates") or [{}])[0]
    text = "".join(p.get("text", "") for p in (candidate.get("content") or {}).get("parts", []))
    usage = result.get("usageMetadata", {})

    # Gemini-2.5 image token-cost: input $1.25/M, output $5/M (Pro); $0.10/$0.40 (Flash)
    in_price = 1.25 if "pro" in model else 0.10
    out_price = 5.00 if "pro" in model else 0.40
    cost = (usage.get("promptTokenCount", 0) * in_price
             + usage.get("candidatesTokenCount", 0) * out_price) / 1_000_000

    out_dict = {
        "text": text, "model": model, "latency_ms": latency_ms,
        "prompt_tokens": usage.get("promptTokenCount", 0),
        "completion_tokens": usage.get("candidatesTokenCount", 0),
        "cost_usd": round(cost, 6),
    }

    # Auto-Trace
    try:
        from aria_llm_trace import log_call, make_hash  # type: ignore
        log_call(
            task_class="vision", worker_model=model, worker_family="google",
            prompt_tokens=out_dict["prompt_tokens"],
            completion_tokens=out_dict["completion_tokens"],
            cost_usd=cost, latency_ms=latency_ms,
            input_hash=make_hash(str(path) + task),
            output_hash=make_hash(text),
            privacy_tier="cloud",
        )
    except Exception:
        pass

    return out_dict


def main() -> int:
    parser = argparse.ArgumentParser(description="Aria Vision Helper via Gemini")
    parser.add_argument("image", help="path to image file")
    parser.add_argument("--task", default="Describe this image in detail. If text is visible, transcribe it.")
    parser.add_argument("--pro", action="store_true", help="use gemini-2.5-pro (default: flash)")
    parser.add_argument("--max-tokens", type=int, default=1500)
    parser.add_argument("--json", action="store_true")
    args = parser.parse_args()

    model = "gemini-2.5-pro" if args.pro else "gemini-2.5-flash"
    out = analyze_image(Path(args.image), task=args.task, model=model, max_tokens=args.max_tokens)
    if args.json:
        print(json.dumps(out, indent=2, ensure_ascii=False))
    else:
        print(out["text"])
        sys.stderr.write(
            f"[vision] {model} · tokens {out['prompt_tokens']}/{out['completion_tokens']} "
            f"· cost ${out['cost_usd']:.6f} · {out['latency_ms']}ms\n"
        )
    return 0


if __name__ == "__main__":
    sys.exit(main())