#!/usr/bin/env python3
"""VLM helper — MULTI-PROVIDER com FALLBACK automático. Judges visuais do editor.

Providers (ambos com VISÃO): kimi (Moonshot, OpenAI-compat) e claude (Anthropic messages).
Qwen entra no futuro (mesma interface, só p/ produção). Trocar/priorizar via env:
    VLM_PROVIDER=claude          # só Claude
    VLM_PROVIDER=claude,kimi     # Claude primário, Kimi fallback
    (default: kimi,claude — Kimi primário barato, Claude fallback)
Fallback é AUTOMÁTICO: se o provider falha (429/sem-chave/erro), tenta o próximo na ordem.
Chaves: .kimi_key / .anthropic_key (gitignored) OU env KIMI_API_KEY / ANTHROPIC_API_KEY.
Uso (inalterado p/ os callers): from vlm import ask_image; ask_image(jpg, prompt) -> (texto, usage)."""
import base64, json, os, sys, time, threading, urllib.request
from concurrent.futures import ThreadPoolExecutor

_DIR = os.path.dirname(__file__)


def _load_key(env_var, fname):
    if os.environ.get(env_var):
        return os.environ[env_var]
    # Procura em (1) scripts/ (compat palmier), (2) root do projeto (padrão VideoTask,
    # gitignored no root .gitignore). Primeiro achado ganha.
    for p in (os.path.join(_DIR, fname), os.path.join(os.path.dirname(_DIR), fname)):
        if os.path.exists(p):
            return open(p).read().strip()
    return None


# ---- providers ----------------------------------------------------------------
# Kimi Tier1: RPM 200, TPM 2M, concurrency 50, TPD ilimitado. throttle por-provider (s entre
# chamadas do MESMO provider). 200 RPM = 0.3s/req; 0.4s dá margem p/ o 429 recuar via backoff.
# (Tier0 era 4.0 por causa do RPM 20 + TPD 1.5M/dia — não é mais o gargalo.)
_MIN_INTERVAL = {"kimi": 0.4, "claude": 1.2}
_LAST = {}
_THROTTLE_LOCK = threading.Lock()


def _throttle(prov):
    # Thread-safe: sob threads, o lock serializa só o ESPAÇAMENTO de dispatch (1 req a cada `iv`s),
    # não a chamada HTTP (que roda concorrente após soltar o lock). Tier1 Kimi: 0.4s = 150 RPM < 200,
    # com até 50 em voo → latência escondida. Sem isso o pool estoura o RPM.
    iv = _MIN_INTERVAL.get(prov, 1.0)
    with _THROTTLE_LOCK:
        dt = time.time() - _LAST.get(prov, 0.0)
        if dt < iv:
            time.sleep(iv - dt)
        _LAST[prov] = time.time()


def _kimi_call(img_b64, prompt, key, temperature):
    body = json.dumps({
        "model": os.environ.get("KIMI_MODEL", "moonshot-v1-8k-vision-preview"),
        "temperature": temperature, "max_tokens": 1024,  # sem isso o rate-limit estima o request GRANDE (429 na borda da cota)
        "messages": [{"role": "user", "content": [
            {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{img_b64}"}},
            {"type": "text", "text": prompt}]}]}).encode()
    hdr = {"Content-Type": "application/json", "Authorization": f"Bearer {key}"}
    r = json.load(urllib.request.urlopen(
        urllib.request.Request("https://api.moonshot.ai/v1/chat/completions", body, hdr), timeout=90))
    return r["choices"][0]["message"]["content"], r.get("usage", {})


def _claude_call(img_b64, prompt, key, temperature):
    body = json.dumps({
        "model": os.environ.get("CLAUDE_VLM_MODEL", "claude-opus-4-8"),
        "max_tokens": 1024, "temperature": temperature,
        "messages": [{"role": "user", "content": [
            {"type": "image", "source": {"type": "base64", "media_type": "image/jpeg", "data": img_b64}},
            {"type": "text", "text": prompt}]}]}).encode()
    hdr = {"Content-Type": "application/json", "x-api-key": key, "anthropic-version": "2023-06-01"}
    r = json.load(urllib.request.urlopen(
        urllib.request.Request("https://api.anthropic.com/v1/messages", body, hdr), timeout=90))
    return "".join(b.get("text", "") for b in r.get("content", []) if b.get("type") == "text"), r.get("usage", {})


PROVIDERS = {
    "kimi":   {"call": _kimi_call,   "key": lambda: _load_key("KIMI_API_KEY", ".kimi_key")},
    "claude": {"call": _claude_call, "key": lambda: _load_key("ANTHROPIC_API_KEY", ".anthropic_key")},
    # "qwen": {...}  # futuro (produção), mesma interface
}


def provider_order():
    env = os.environ.get("VLM_PROVIDER", "").strip()
    if env:
        order = [p.strip() for p in env.split(",") if p.strip() in PROVIDERS]
        if order:
            return order
    return ["kimi", "claude"]  # default: Kimi barato primário, Claude fallback


def _ledger_usage(kind, provider, usage, elapsed_s, n_images=1, ok=True, err=None):
    """Append-only vlm_usage.jsonl no cwd (workdir do cut) p/ custo/tempo pós-run.
    Ativo se VIDEOTASK_VLM_LEDGER não for 0 (default: ON)."""
    if os.environ.get("VIDEOTASK_VLM_LEDGER", "1").strip() in ("0", "false", "no"):
        return
    try:
        import json as _json
        from datetime import datetime, timezone
        path = os.environ.get("VIDEOTASK_VLM_LEDGER_PATH") or os.path.join(os.getcwd(), "vlm_usage.jsonl")
        # rough USD: Moonshot vision ~$0.002/1k in + $0.008? — documentado como ESTIMATE
        # defaults conservadores (ajustar via env se quiser)
        pin = float(os.environ.get("VLM_USD_PER_1K_INPUT", "0.00015"))   # kimi-ish cheap
        pout = float(os.environ.get("VLM_USD_PER_1K_OUTPUT", "0.0025"))
        if provider == "claude":
            pin = float(os.environ.get("VLM_CLAUDE_USD_PER_1K_INPUT", "0.003"))
            pout = float(os.environ.get("VLM_CLAUDE_USD_PER_1K_OUTPUT", "0.015"))
        u = usage or {}
        tin = int(u.get("prompt_tokens") or u.get("input_tokens") or 0)
        tout = int(u.get("completion_tokens") or u.get("output_tokens") or 0)
        # anthropic usage shape
        if not tin and isinstance(u.get("input_tokens"), int):
            tin = u["input_tokens"]
        cost = (tin / 1000.0) * pin + (tout / 1000.0) * pout
        row = {
            "ts": datetime.now(timezone.utc).isoformat(),
            "kind": kind,
            "provider": provider,
            "ok": ok,
            "n_images": n_images,
            "prompt_tokens": tin,
            "completion_tokens": tout,
            "elapsed_s": round(float(elapsed_s), 3),
            "cost_usd_est": round(cost, 6),
            "err": err,
        }
        with open(path, "a") as f:
            f.write(_json.dumps(row, ensure_ascii=False) + "\n")
    except Exception:
        pass


def ask_image(jpg, prompt, temperature=0.1, max_retries=4):
    """Tenta cada provider na ordem; retry por-provider (429/erro transiente); cai pro próximo
    se esgotar/sem-chave. Retorna (texto, usage). ERR:<detalhe> só se TODOS falharem."""
    img = base64.b64encode(open(jpg, "rb").read()).decode()
    errs = []
    for prov in provider_order():
        key = PROVIDERS[prov]["key"]()
        if not key:
            errs.append(f"{prov}:sem-chave")
            continue
        call = PROVIDERS[prov]["call"]
        for i in range(max_retries):
            try:
                _throttle(prov)
                t0 = time.time()
                text, usage = call(img, prompt, key, temperature)
                _ledger_usage("ask_image", prov, usage, time.time() - t0, n_images=1, ok=True)
                return text, usage
            except Exception as e:
                code = getattr(e, "code", None)
                detail = type(e).__name__
                retry_after = None
                if code is not None:
                    try:
                        detail = f"HTTP{code}:{e.read().decode()[:160]}"
                    except Exception:
                        detail = f"HTTP{code}"
                    try:
                        retry_after = float(e.headers.get("Retry-After") or 0) or None
                    except Exception:
                        pass
                print(f"  [vlm:{prov}] tentativa {i+1}/{max_retries}: {detail}", file=sys.stderr)
                is_daily = code == 429 and "TPD" in detail  # cota diária: adianta pro próximo provider
                if is_daily or i == max_retries - 1:
                    errs.append(f"{prov}:{detail[:60]}")
                    break
                time.sleep(retry_after if retry_after else (8 * (i + 1) if code == 429 else 5 * (i + 1)))
    return f"ERR:{';'.join(errs)}", {}


def ask_images(jpgs, prompt, temperature=0.1):
    """Várias imagens numa chamada (grid) — corta chamadas/tokens. Mesmo fallback multi-provider."""
    b64s = [base64.b64encode(open(j, "rb").read()).decode() for j in jpgs]
    errs = []
    for prov in provider_order():
        key = PROVIDERS[prov]["key"]()
        if not key:
            errs.append(f"{prov}:sem-chave"); continue
        try:
            _throttle(prov)
            t0 = time.time()
            if prov == "kimi":
                content = [{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{b}"}} for b in b64s]
                content.append({"type": "text", "text": prompt})
                body = json.dumps({"model": os.environ.get("KIMI_MODEL", "moonshot-v1-8k-vision-preview"),
                                   "temperature": temperature,
                                   "messages": [{"role": "user", "content": content}]}).encode()
                hdr = {"Content-Type": "application/json", "Authorization": f"Bearer {key}"}
                r = json.load(urllib.request.urlopen(
                    urllib.request.Request("https://api.moonshot.ai/v1/chat/completions", body, hdr), timeout=120))
                text, usage = r["choices"][0]["message"]["content"], r.get("usage", {})
                _ledger_usage("ask_images", prov, usage, time.time() - t0, n_images=len(b64s), ok=True)
                return text, usage
            else:  # claude
                content = [{"type": "image", "source": {"type": "base64", "media_type": "image/jpeg", "data": b}} for b in b64s]
                content.append({"type": "text", "text": prompt})
                body = json.dumps({"model": os.environ.get("CLAUDE_VLM_MODEL", "claude-opus-4-8"),
                                   "max_tokens": 2048, "temperature": temperature,
                                   "messages": [{"role": "user", "content": content}]}).encode()
                hdr = {"Content-Type": "application/json", "x-api-key": key, "anthropic-version": "2023-06-01"}
                r = json.load(urllib.request.urlopen(
                    urllib.request.Request("https://api.anthropic.com/v1/messages", body, hdr), timeout=120))
                text = "".join(b.get("text", "") for b in r.get("content", []) if b.get("type") == "text")
                usage = r.get("usage", {})
                _ledger_usage("ask_images", prov, usage, time.time() - t0, n_images=len(b64s), ok=True)
                return text, usage
        except Exception as e:
            detail = f"HTTP{getattr(e,'code','')}" if hasattr(e, "code") else type(e).__name__
            print(f"  [vlm:{prov}] batch falhou: {detail}", file=sys.stderr)
            errs.append(f"{prov}:{detail}")
    return f"ERR:{';'.join(errs)}", {}


PAR_WORKERS = int(os.environ.get("VLM_WORKERS", "10"))


def par_map(fn, items, workers=None):
    """Roda fn(item) concorrente preservando ORDEM. O throttle global (thread-safe) espaça o dispatch
    ao RPM do tier; os workers só escondem a latência de rede/inferência. Exceção vira o retorno de fn
    (não derruba o batch) — deixe fn tratar erro e devolver default. workers=1 => serial (debug)."""
    items = list(items)
    if not items:
        return []
    w = workers or PAR_WORKERS
    if w <= 1:
        return [fn(x) for x in items]
    with ThreadPoolExecutor(max_workers=min(w, len(items))) as ex:
        return list(ex.map(fn, items))


def key():  # back-comp
    return PROVIDERS["kimi"]["key"]()
