"""Optional AI image for inserts, through OpenRouter (Contract 8).

The key comes from OPENROUTER_API_KEY in the environment or from `.env` in the skill root.
Without a key the helper exits 2 and explains that this step is optional (fetch_images.py finds
free photos with no key).

The model is picked from OpenRouter's public model list (models whose output includes images),
preferring a fast, cheap, current one; `--model` forces another. The image is resized to cover
`--size` and saved to <edit>/assets/, with an entry in <edit>/assets/credits.json.

Before decoding, the answer (image link or data) is saved to <edit>/assets/.ia-ultima-resposta.json.
If saving the image fails after a paid generation (link expired, network, bad data), the error
says so and `--reuse-last` tries again from that file without calling the model (no new cost).

Usage:
    uv run python helpers/gen_image.py "<prompt>" --edit-dir DIR [--size 1080x900] [-o name.png]
        [--model ID] [--raw] [--dry-run] [--reuse-last]
"""

from __future__ import annotations

import argparse
import base64
import binascii
import hashlib
import io
import json
import re
import sys
import unicodedata
from pathlib import Path

import requests
from PIL import Image

sys.path.insert(0, str(Path(__file__).resolve().parent))
from common import load_json, now_iso, save_json  # noqa: E402
from fetch_images import EnvFileError, read_env_key  # noqa: E402  (one .env reader for both helpers)

USER_AGENT = "editor-video/0.1 (suporte@sistemajarvis.com)"
MODELS_URL = "https://openrouter.ai/api/v1/models"
CHAT_URL = "https://openrouter.ai/api/v1/chat/completions"
# Preferred order: current, fast and cheap first. The live list decides what exists.
PREFERRED_MODELS = (
    "google/gemini-3.1-flash-image",
    "google/gemini-3.1-flash-image-preview",
    "google/gemini-2.5-flash-image",
    "openai/gpt-5-image-mini",
    "google/gemini-3-pro-image",
    "openai/gpt-5-image",
)
FALLBACK_MODEL = PREFERRED_MODELS[0]
ASPECTS = ("1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9", "21:9")
PROMPT_SUFFIX = ("Photorealistic, natural light, sharp focus. No text, letters, logos, brand marks "
                 "or watermarks anywhere in the image.")


class GenError(Exception):
    def __init__(self, message: str, code: int = 1):
        super().__init__(message)
        self.code = code


def parse_size(value: str) -> tuple[int, int]:
    m = re.fullmatch(r"\s*(\d{2,5})\s*[xX]\s*(\d{2,5})\s*", value or "")
    if not m:
        raise GenError(f"tamanho inválido: {value!r}. Use LARGURAxALTURA, ex: 1080x900", 2)
    w, h = int(m.group(1)), int(m.group(2))
    if not (64 <= w <= 4096 and 64 <= h <= 4096):
        raise GenError("o tamanho precisa ficar entre 64 e 4096 px de cada lado", 2)
    return w, h


def nearest_aspect(w: int, h: int) -> str:
    target = w / h

    def ratio(a: str) -> float:
        x, y = a.split(":")
        return int(x) / int(y)

    return min(ASPECTS, key=lambda a: abs(ratio(a) - target) / target)


def slugify(text: str, limit: int = 32) -> str:
    t = unicodedata.normalize("NFKD", text).encode("ascii", "ignore").decode("ascii").lower()
    t = re.sub(r"[^a-z0-9]+", "-", t).strip("-")
    return (t[:limit].rstrip("-")) or "imagem"


def image_models(timeout: float = 20) -> list[dict]:
    """Public model list filtered to models that can output images (no key needed)."""
    r = requests.get(MODELS_URL, headers={"User-Agent": USER_AGENT}, timeout=timeout)
    r.raise_for_status()
    out = []
    for m in r.json().get("data", []):
        arch = m.get("architecture") or {}
        if "image" in (arch.get("output_modalities") or []) and not str(m.get("id", "")).startswith("openrouter/"):
            out.append(m)
    return out


def pick_model(requested: str | None) -> tuple[str, str]:
    """-> (model id, note). Uses the live list; falls back to a known id when offline."""
    try:
        models = image_models()
    except (requests.RequestException, ValueError):
        if requested:
            return requested, "lista de modelos indisponível; usando o modelo pedido"
        return FALLBACK_MODEL, "lista de modelos indisponível; usando o padrão"
    ids = {m["id"] for m in models}
    if requested:
        if requested not in ids:
            raise GenError(f"o modelo '{requested}' não gera imagem no OpenRouter agora. Opções: {', '.join(sorted(ids))}", 2)
        return requested, "modelo pedido"
    for mid in PREFERRED_MODELS:
        if mid in ids:
            return mid, "modelo preferido disponível"
    if not models:
        return FALLBACK_MODEL, "nenhum modelo de imagem listado; usando o padrão"

    def price(m: dict) -> float:
        try:
            return float((m.get("pricing") or {}).get("image_output") or 1e9)
        except ValueError:
            return 1e9

    best = sorted(models, key=lambda m: (price(m), -int(m.get("created") or 0)))[0]
    return best["id"], "modelo de imagem mais barato da lista"


def build_request(prompt: str, model: str, size: tuple[int, int], raw: bool) -> dict:
    text = prompt.strip() if raw else f"{prompt.strip()}\n\n{PROMPT_SUFFIX}"
    return {
        "model": model,
        "messages": [{"role": "user", "content": text}],
        "modalities": ["image", "text"],
        "image_config": {"aspect_ratio": nearest_aspect(*size)},
    }


LAST_RESPONSE = ".ia-ultima-resposta.json"


def image_from_url(url: str, saved: str) -> bytes:
    """Bytes of an image the model returned (data: URL or link). `saved` names where the answer is kept."""
    if not isinstance(url, str) or not url:
        raise GenError(f"o modelo devolveu um endereço de imagem vazio. A resposta ficou em {saved}.")
    if url.startswith("data:"):
        if "," not in url:
            raise GenError(f"a imagem veio num formato que não consegui ler. A resposta ficou em {saved}.")
        try:
            return base64.b64decode(url.split(",", 1)[1])
        except (binascii.Error, ValueError):
            raise GenError(f"a imagem veio com dados inválidos (base64). A resposta ficou em {saved}; rode de novo "
                           f"com --reuse-last se quiser tentar outra vez sem pagar.")
    if url.startswith("http"):
        try:
            got = requests.get(url, headers={"User-Agent": USER_AGENT}, timeout=60)
            got.raise_for_status()
            return got.content
        except requests.RequestException as e:
            code = getattr(getattr(e, "response", None), "status_code", None)
            why = f"erro {code}" if code else e.__class__.__name__
            raise GenError(f"a imagem foi gerada, mas não consegui baixar ({why}). O link ficou em {saved}: rode de novo "
                           f"com --reuse-last para tentar baixar sem pagar outra vez.")
    raise GenError(f"o modelo devolveu um endereço de imagem que não reconheço. A resposta ficou em {saved}.")


def request_image(payload: dict, key: str, timeout: float = 180, save_to: Path | None = None) -> bytes:
    headers = {"Authorization": f"Bearer {key}", "Content-Type": "application/json",
               "User-Agent": USER_AGENT, "X-Title": "editor-video"}
    try:
        r = requests.post(CHAT_URL, headers=headers, data=json.dumps(payload), timeout=timeout)
    except requests.RequestException as e:
        raise GenError(f"não consegui falar com o OpenRouter ({e.__class__.__name__}). Confira a internet.")
    if r.status_code == 401:
        raise GenError("o OpenRouter recusou a chave OPENROUTER_API_KEY. Confira se ela está certa.", 2)
    if r.status_code == 402:
        raise GenError("a conta do OpenRouter está sem crédito. Adicione crédito ou use fetch_images.py.")
    if r.status_code == 429:
        raise GenError("limite de uso do OpenRouter atingido. Tente de novo em alguns minutos.")
    if r.status_code >= 400:
        raise GenError(f"o OpenRouter respondeu erro {r.status_code}: {r.text[:300]}")
    try:
        data = r.json()
        message = data["choices"][0]["message"]
    except (ValueError, KeyError, IndexError, TypeError):
        raise GenError("resposta inesperada do OpenRouter (sem mensagem).")
    if not isinstance(message, dict):
        raise GenError("resposta inesperada do OpenRouter (mensagem vazia). Nada foi salvo; tente de novo.")
    urls = []
    for img in message.get("images") or []:
        url = ((img or {}).get("image_url") or {}).get("url", "") if isinstance(img, dict) else ""
        if isinstance(url, str) and url.startswith(("data:", "http")):
            urls.append(url)
    if not urls:
        raise GenError("o modelo não devolveu imagem. Tente reescrever o pedido de forma mais visual.")
    saved = "a resposta"
    if save_to is not None:
        try:
            save_json(save_to, {"model": payload.get("model"), "urls": urls, "saved_at": now_iso()})
            saved = f"assets/{save_to.name}"
        except OSError:
            pass
    return image_from_url(urls[0], saved)


def cover(img: Image.Image, size: tuple[int, int]) -> Image.Image:
    w, h = size
    s = max(w / img.width, h / img.height)
    nw, nh = max(1, round(img.width * s)), max(1, round(img.height * s))
    r = img.convert("RGB").resize((nw, nh), Image.LANCZOS)
    x, y = (nw - w) // 2, (nh - h) // 2
    return r.crop((x, y, x + w, y + h))


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser(description="Gera uma imagem com IA pelo OpenRouter (opcional) para usar nos inserts.")
    ap.add_argument("prompt", help="descrição da imagem")
    ap.add_argument("--edit-dir", required=True, help="pasta edit da sessão")
    ap.add_argument("--size", default="1080x900", help="tamanho final LARGURAxALTURA (padrão 1080x900)")
    ap.add_argument("-o", "--output", help="nome do arquivo em assets/ (padrão: gerado do texto)")
    ap.add_argument("--model", help="força um modelo do OpenRouter")
    ap.add_argument("--raw", action="store_true", help="manda o texto exatamente como está (sem o complemento sem texto/logo)")
    ap.add_argument("--dry-run", action="store_true", help="só mostra o pedido que seria enviado, sem gerar")
    ap.add_argument("--reuse-last", action="store_true",
                    help="não gera de novo: salva a imagem da última resposta guardada (depois de um erro ao baixar)")
    args = ap.parse_args(argv)

    try:
        if not args.prompt.strip():
            raise GenError("o pedido está vazio", 2)
        size = parse_size(args.size)
        edit = Path(args.edit_dir).resolve()
        if not edit.is_dir():
            raise GenError(f"pasta edit não encontrada: {edit}", 2)
        try:
            key = read_env_key("OPENROUTER_API_KEY")
        except EnvFileError as e:
            raise GenError(str(e), 2)
        last_path = edit / "assets" / LAST_RESPONSE
        if args.reuse_last:
            last = load_json(last_path) if last_path.exists() else None
            if not isinstance(last, dict) or not last.get("urls"):
                raise GenError("não há resposta guardada para reaproveitar (assets/.ia-ultima-resposta.json). "
                               "Gere a imagem normalmente.", 2)
            key = key or "reuse"
        if not key and not args.dry_run:
            raise GenError(
                "OPENROUTER_API_KEY não configurada. Gerar imagem com IA é opcional: dá para seguir sem ela "
                "usando fetch_images.py (fotos livres). Para ativar, crie uma chave em openrouter.ai e coloque "
                "OPENROUTER_API_KEY=... no ambiente ou no arquivo .env da skill.", 2)

        if args.reuse_last:
            model, note = str(last.get("model") or args.model or FALLBACK_MODEL), "resposta guardada"
        else:
            model, note = pick_model(args.model)
        payload = build_request(args.prompt, model, size, args.raw)
        name = args.output or f"ia-{slugify(args.prompt)}-{hashlib.sha1((args.prompt + model).encode()).hexdigest()[:6]}.png"
        name = Path(name).name
        if not name.lower().endswith((".png", ".jpg", ".jpeg")):
            name += ".png"
        out_path = edit / "assets" / name

        if args.dry_run:
            print(f"modelo: {model} ({note})")
            print(f"arquivo: {out_path}")
            print("pedido (não enviado):")
            print(json.dumps({"url": CHAT_URL, "payload": payload, "key": "configurada" if key else "ausente"},
                             ensure_ascii=False, indent=2))
            return 0

        if args.reuse_last:
            print(f"usando a última resposta guardada ({model}), sem gerar de novo...")
            raw = image_from_url(str(last["urls"][0]), f"assets/{LAST_RESPONSE}")
        else:
            print(f"gerando com {model}...")
            (edit / "assets").mkdir(parents=True, exist_ok=True)
            raw = request_image(payload, key or "", save_to=last_path)
        try:
            img = Image.open(io.BytesIO(raw))
            img.load()
        except Exception:
            raise GenError(f"o OpenRouter devolveu um arquivo que não é imagem. A resposta ficou em assets/{LAST_RESPONSE}.")
        final = cover(img, size)
        out_path.parent.mkdir(parents=True, exist_ok=True)
        tmp = out_path.with_name(".tmp-" + out_path.name)
        if name.lower().endswith(".png"):
            final.save(tmp, "PNG")
        else:
            final.save(tmp, "JPEG", quality=93)
        tmp.replace(out_path)

        credits_path = edit / "assets" / "credits.json"
        credits = load_json(credits_path, default=[])
        if not isinstance(credits, list):
            credits = []
        credits.append({
            "file": f"assets/{name}", "source": "openrouter", "model": model, "prompt": args.prompt,
            "url": None, "author": f"imagem gerada por IA ({model})", "license": "gerada por IA",
            "width": size[0], "height": size[1], "added_at": now_iso(),
        })
        save_json(credits_path, credits)
        print(f"ok: assets/{name}  {size[0]}x{size[1]}  modelo {model}")
        return 0
    except GenError as e:
        print(f"erro: {e}", file=sys.stderr)
        return e.code


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