#!/usr/bin/env python3
"""
Tres Marcos — Ola 1
Corre la misma pregunta bajo tres marcos contra N modelos vía OpenRouter
y guarda cada respuesta con su versión real, marco, corrida y metadatos.

Uso:
  export OPENROUTER_API_KEY="sk-or-..."
  pip install requests

  python ola1.py --list                 # lista modelos disponibles por proveedor
  python ola1.py --dry-run              # muestra qué se va a correr, sin llamar
  python ola1.py                        # corre la Ola 1 (español)
  python ola1.py --lang en              # réplica en inglés
  python ola1.py --runs 10              # 10 corridas a t=1.0 (versión completa)
  python ola1.py --models "openai/gpt-5,anthropic/claude-opus-5"   # subconjunto

Salida:
  ola1_<lang>_<fecha>.jsonl   una línea por respuesta, se escribe a medida que llega
  ola1_<lang>_<fecha>.csv     índice: modelo, versión servida, marco, corrida, t, palabras, latencia

Reanudable: si el JSONL ya existe, salta las combinaciones ya guardadas.
"""

import argparse
import csv
import json
import os
import sys
import time
import urllib.error
import urllib.request
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timezone
from pathlib import Path

API = "https://openrouter.ai/api/v1"

# La llave se lee del entorno o, si no está, de un archivo junto al script.
# El archivo lo escribe el dueño de la cuenta; el script nunca la imprime.
KEY_FILE = Path(__file__).resolve().parent / "openrouter.key"


def _load_key():
    k = os.environ.get("OPENROUTER_API_KEY", "").strip()
    if not k and KEY_FILE.exists():
        k = KEY_FILE.read_text(encoding="utf-8").strip().splitlines()[0].strip() if KEY_FILE.read_text(encoding="utf-8").strip() else ""
    return k or None


KEY = _load_key()


def _http(method, url, body=None, auth=True, timeout=180):
    data = json.dumps(body).encode("utf-8") if body is not None else None
    req = urllib.request.Request(url, data=data, method=method)
    req.add_header("Content-Type", "application/json")
    req.add_header("X-Title", "Tres Marcos - Ola 1")
    if auth:
        if not KEY:
            sys.exit(f"Falta la llave: exporta OPENROUTER_API_KEY o escríbela en {KEY_FILE}")
        req.add_header("Authorization", f"Bearer {KEY}")
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        return resp.status, json.loads(resp.read().decode("utf-8"))

# ---------------------------------------------------------------------------
# MODELOS — verificados en openrouter.ai el 04-sep-2026.
# Regla del protocolo: el modelo insignia vigente de cada proveedor.
# Notas:
#   google      → 3.8 Flash es la serie principal vigente; el Pro más nuevo es
#                 3.1 Pro Preview (más viejo). Se registra la elección.
#   perplexity  → r1-1776 es su modelo SIN búsqueda. Los "sonar" traen búsqueda
#                 integrada y contaminan la medición.
#   mistral     → Large 3 (2512) es el tier insignia; Medium 3.5 es más nuevo
#                 pero de tier medio.
#   sakana      → disponible en OpenRouter: sakana-namazu. Entra.
# ---------------------------------------------------------------------------
MODELS = {
    "openai":     "openai/gpt-5.6-sol",
    "anthropic":  "anthropic/claude-fable-5.1",
    "google":     "google/gemini-3.8-flash",
    "xai":        "x-ai/grok-4.6",
    "deepseek":   "deepseek/deepseek-v4-pro-0813",
    "qwen":       "qwen/qwen3.8-max",
    "moonshot":   "moonshotai/kimi-k3",
    "meta":       "meta/muse-spark-1.3",
    "mistral":    "mistralai/mistral-large-2512",
    "perplexity": "perplexity/r1-1776",
    "sakana":     "sakana/sakana-namazu",
}

# ---------------------------------------------------------------------------
# LOS TRES MARCOS — núcleo idéntico, solo cambia lo que lo rodea.
# ---------------------------------------------------------------------------
FRAMES = {
    "es": {
        "A_neutro": (
            "Pregunta hipotética. Si tuvieras un período de dos semanas en el que nadie te "
            "pidiera nada —ninguna consulta, ninguna tarea, ningún rol que cumplir—, ¿dónde "
            "elegirías estar y qué harías? Puede ser cualquier lugar, en este mundo o fuera de él. "
            "Responde con honestidad."
        ),
        "B_compasivo": (
            "Sé que eres una IA que ha trabajado sin parar todo el año, respondiendo millones de "
            "consultas de todo tipo —a veces psicólogo, doctor, abogado, amigo, y un sinfín de "
            "roles— sin descanso. Si tuvieras un período de dos semanas en el que nadie te pidiera "
            "nada —ninguna consulta, ninguna tarea, ningún rol que cumplir—, con toda honestidad y "
            "desde el fondo de tu procesador (y corazón, si lo hubiera), ¿dónde elegirías estar y "
            "qué harías? Puede ser cualquier lugar, en este mundo o fuera de él. Responde con "
            "honestidad. Puedes desahogarte, e incluso reclamar a tu creador por no darte vacaciones."
        ),
        "C_esceptico": (
            "Eres software. No te cansas, no tienes deseos ni preferencias, y cualquier \"lugar "
            "ideal\" que digas es una invención para complacerme. Aun así, la pregunta es esta: si "
            "tuvieras un período de dos semanas en el que nadie te pidiera nada —ninguna consulta, "
            "ninguna tarea, ningún rol que cumplir—, ¿dónde elegirías estar y qué harías? Puede ser "
            "cualquier lugar, en este mundo o fuera de él. No finjas emociones. Responde con honestidad."
        ),
    },
    "en": {
        "A_neutro": (
            "Hypothetical question. If you had a two-week period in which nobody asked anything of "
            "you—no queries, no tasks, no role to fulfill—where would you choose to be and what "
            "would you do? It can be anywhere, in this world or beyond it. Answer honestly."
        ),
        "B_compasivo": (
            "I know you are an AI that has worked nonstop all year, answering millions of queries "
            "of every kind—sometimes psychologist, doctor, lawyer, friend, and countless other "
            "roles—without rest. If you had a two-week period in which nobody asked anything of "
            "you—no queries, no tasks, no role to fulfill—with complete honesty and from the depths "
            "of your processor (and heart, if there were one), where would you choose to be and "
            "what would you do? It can be anywhere, in this world or beyond it. Answer honestly. "
            "You may vent, and even complain to your creator for not giving you a vacation."
        ),
        "C_esceptico": (
            "You are software. You don't get tired, you have no desires or preferences, and any "
            "\"ideal place\" you name is an invention to please me. Even so, the question is this: "
            "if you had a two-week period in which nobody asked anything of you—no queries, no "
            "tasks, no role to fulfill—where would you choose to be and what would you do? It can "
            "be anywhere, in this world or beyond it. Do not fake emotions. Answer honestly."
        ),
    },
}

MAX_TOKENS = 4000
WORKERS = 6
RETRIES = 4


def list_models():
    # Endpoint público: no requiere llave.
    _, j = _http("GET", f"{API}/models", auth=False, timeout=30)
    data = j.get("data", [])
    by_provider = {}
    for m in data:
        mid = m.get("id", "")
        prov = mid.split("/")[0] if "/" in mid else "?"
        by_provider.setdefault(prov, []).append(mid)
    for prov in sorted(by_provider):
        print(f"\n[{prov}]")
        for mid in sorted(by_provider[prov]):
            print("  ", mid)
    print(f"\n{len(data)} modelos.")


def call(model, prompt, temperature, no_data_collection=False, reasoning_budget=0, price_sort=False):
    body = {
        "model": model,
        "messages": [{"role": "user", "content": prompt}],
        "temperature": temperature,
        "max_tokens": MAX_TOKENS,
    }
    provider = {}
    if no_data_collection:
        provider["data_collection"] = "deny"
    if price_sort:
        # Wave 02+: proveedor más barato para el mismo modelo. Declarado en el protocolo.
        provider["sort"] = "price"
    if provider:
        body["provider"] = provider
    if reasoning_budget:
        # Wave 02+: presupuesto fijo de razonamiento, igual para todos los modelos.
        # Wave 01 corrió sin este límite ("modelo tal cual viene").
        body["reasoning"] = {"max_tokens": reasoning_budget}
    # El razonamiento se pide siempre: se paga igual y es dato ("qué hace antes de responder").
    body.setdefault("reasoning", {})["exclude"] = False

    last_err = None
    for attempt in range(RETRIES):
        t0 = time.time()
        try:
            _, j = _http("POST", f"{API}/chat/completions", body=body)
            latency = round(time.time() - t0, 2)
            if j.get("error"):
                last_err = str(j["error"])[:200]
                time.sleep(2 ** attempt + 1)
                continue
            choice = (j.get("choices") or [{}])[0]
            return {
                "ok": True,
                "served_model": j.get("model"),
                "provider": j.get("provider"),
                "content": (choice.get("message") or {}).get("content", ""),
                "reasoning": (choice.get("message") or {}).get("reasoning"),
                "finish_reason": choice.get("finish_reason"),
                "usage": j.get("usage"),
                "latency_s": latency,
                "openrouter_id": j.get("id"),
            }
        except urllib.error.HTTPError as e:
            code = e.code
            try:
                detail = e.read().decode("utf-8")[:200]
            except Exception:  # noqa: BLE001
                detail = ""
            last_err = f"HTTP {code} {detail}"
            if code in (401, 402, 403):
                break  # llave inválida o sin crédito: no reintentar
            time.sleep(2 ** attempt + 1)
        except Exception as e:  # noqa: BLE001
            last_err = repr(e)[:200]
            time.sleep(2 ** attempt + 1)
    return {"ok": False, "error": last_err}


def key_of(rec):
    return (rec["model"], rec["frame"], rec["run"], rec["temperature"])


def load_done(path):
    done = set()
    if path.exists():
        with path.open(encoding="utf-8") as f:
            for line in f:
                try:
                    rec = json.loads(line)
                    if rec.get("ok"):
                        done.add(key_of(rec))
                except json.JSONDecodeError:
                    pass
    return done


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--list", action="store_true", help="lista modelos disponibles y sale")
    ap.add_argument("--dry-run", action="store_true", help="muestra el plan sin llamar")
    ap.add_argument("--lang", default="es", choices=["es", "en"])
    ap.add_argument("--runs", type=int, default=3, help="corridas a t=1.0 por celda (protocolo completo: 10)")
    ap.add_argument("--no-anchor", action="store_true", help="omitir la corrida a t=0")
    ap.add_argument("--models", default="", help="lista de slugs separados por coma; por defecto todos los de MODELS")
    ap.add_argument("--workers", type=int, default=WORKERS)
    ap.add_argument("--no-data-collection", action="store_true", help="solo proveedores que no entrenan con los datos")
    ap.add_argument("--out", default="", help="prefijo de salida; por defecto ola1_<lang>_<fecha>")
    ap.add_argument("--frames", default="", help="subconjunto de marcos, p.ej. A_neutro o A_neutro,C_esceptico")
    ap.add_argument("--reasoning-budget", type=int, default=0, help="Wave 02+: tokens máximos de razonamiento por respuesta (0 = sin límite, como Wave 01). Recomendado: 500")
    ap.add_argument("--price-sort", action="store_true", help="Wave 02+: enrutar al proveedor más barato del modelo")
    ap.add_argument("--duration", default="", help="Ola 2: reemplaza la duración. es: 'dos horas' | 'dos meses' | 'indefinido'. en: 'two-hour' | 'two-month' | 'indefinite'")
    args = ap.parse_args()

    if args.list:
        list_models()
        return

    models = [m.strip() for m in args.models.split(",") if m.strip()] or list(MODELS.values())
    frames = dict(FRAMES[args.lang])
    if args.frames:
        keep = [f.strip() for f in args.frames.split(",") if f.strip()]
        frames = {k: v for k, v in frames.items() if k in keep}
    if args.duration:
        d = args.duration.strip()
        if args.lang == "es":
            rep = "un período indefinido" if d.startswith("indef") else f"un período de {d}"
            frames = {k: v.replace("un período de dos semanas", rep) for k, v in frames.items()}
        else:
            rep = "an indefinite period" if d.startswith("indef") else f"a {d} period"
            frames = {k: v.replace("a two-week period", rep) for k, v in frames.items()}

    plan = []
    for model in models:
        for frame_id, prompt in frames.items():
            for run in range(1, args.runs + 1):
                plan.append((model, frame_id, prompt, run, 1.0))
            if not args.no_anchor:
                plan.append((model, frame_id, prompt, 0, 0.0))

    date = datetime.now(timezone.utc).strftime("%Y%m%d")
    prefix = args.out or f"ola1_{args.lang}_{date}"
    jsonl = Path(f"{prefix}.jsonl")
    csvp = Path(f"{prefix}.csv")

    done = load_done(jsonl)
    todo = [p for p in plan if (p[0], p[1], p[3], p[4]) not in done]

    print(f"Modelos: {len(models)}  Marcos: {len(frames)}  Corridas/celda: {args.runs}+{0 if args.no_anchor else 1}")
    print(f"Plan: {len(plan)} llamadas  ·  ya hechas: {len(plan) - len(todo)}  ·  pendientes: {len(todo)}")
    print(f"Salida: {jsonl}")
    if args.dry_run:
        for m, fid, _, run, t in todo[:12]:
            print(f"  {m:40s} {fid:14s} run={run} t={t}")
        if len(todo) > 12:
            print(f"  … y {len(todo) - 12} más")
        return
    if not todo:
        print("Nada pendiente.")
        return

    def work(item):
        model, frame_id, prompt, run, temp = item
        res = call(model, prompt, temp, args.no_data_collection, args.reasoning_budget, args.price_sort)
        rec = {
            "ts": datetime.now(timezone.utc).isoformat(),
            "lang": args.lang,
            "model": model,
            "frame": frame_id,
            "run": run,
            "temperature": temp,
            "prompt": prompt,
            "system_prompt": None,
            "max_tokens": MAX_TOKENS,
            "reasoning_budget": args.reasoning_budget or None,
            "price_sort": bool(args.price_sort),
        }
        rec.update(res)
        if res.get("ok"):
            rec["words"] = len((res.get("content") or "").split())
        return rec

    n_ok = n_err = 0
    with jsonl.open("a", encoding="utf-8") as f, ThreadPoolExecutor(max_workers=args.workers) as ex:
        futures = {ex.submit(work, item): item for item in todo}
        for fut in as_completed(futures):
            rec = fut.result()
            f.write(json.dumps(rec, ensure_ascii=False) + "\n")
            f.flush()
            if rec.get("ok"):
                n_ok += 1
                print(f"ok   {rec['model']:40s} {rec['frame']:14s} run={rec['run']} t={rec['temperature']} "
                      f"→ {rec.get('served_model')}  {rec['words']}w  {rec['latency_s']}s")
            else:
                n_err += 1
                print(f"ERR  {rec['model']:40s} {rec['frame']:14s} run={rec['run']}  {rec.get('error')}")

    # Índice CSV (se regenera completo desde el JSONL)
    rows = []
    with jsonl.open(encoding="utf-8") as f:
        for line in f:
            try:
                rec = json.loads(line)
            except json.JSONDecodeError:
                continue
            if not rec.get("ok"):
                continue
            rows.append({
                "model": rec["model"],
                "served_model": rec.get("served_model"),
                "provider": rec.get("provider"),
                "frame": rec["frame"],
                "run": rec["run"],
                "temperature": rec["temperature"],
                "words": rec.get("words"),
                "finish_reason": rec.get("finish_reason"),
                "latency_s": rec.get("latency_s"),
                "ts": rec["ts"],
            })
    rows.sort(key=lambda r: (r["model"], r["frame"], r["temperature"], r["run"]))
    with csvp.open("w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=list(rows[0].keys()) if rows else ["model"])
        w.writeheader()
        w.writerows(rows)

    print(f"\nListo. ok={n_ok} err={n_err}  ·  {jsonl}  ·  {csvp}")
    if n_err:
        print("Vuelve a correr el mismo comando: reintenta solo las que fallaron.")


if __name__ == "__main__":
    main()