#!/usr/bin/env python3
"""Una CAPTURA de la barra lateral de IA contestando, pintada por un compositor de verdad.

    python3 scripts/atuq-captura-ia.py --question "¿qué es una distribución de GNU/Linux?"
    python3 scripts/atuq-captura-ia.py --archive "el gato"     # la otra pregunta: a tus páginas

No es un guardián: es evidencia para una persona. Los guardianes del §6.7 y del §6.3 miden por el
`dump` —qué dijo la extensión, qué contestó el host, qué quedó en disco— y eso es lo correcto para
un veredicto automático, pero no puede mostrar lo único que el usuario ve: **que el panel se pinta
y que la respuesta está ahí**. Por eso esto corre bajo `sway headless` y termina en un PNG.

Es el mismo arnés que usa `test-atuq-descargas.py` (sway con `WLR_BACKENDS=headless`, un `atuq`
adentro), más `grim` para la foto — que el cierre de `escritorio-sway` ya trae.

⚠ Dos corridas, y la primera no es desperdicio: Gecko le asigna a cada extensión una URL base por
PERFIL, así que hay que preguntarle primero cuál le tocó. Adivinar un UUID sería inventar.

REQUISITOS: rootfs de sway (`hydrate-profile.py escritorio-sway`), y `atuq`, `puriy-costura`,
`llama-cpp` y los modelos sellados.
"""
import argparse
import os
import shutil
import subprocess
import sys
import tempfile

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
RFS = os.environ.get("RFS", "/mnt/cosecha/escritorios/sway-rootfs")
ALPINE = os.path.join(ROOT, ".dev-fs/alpine")
HAMMER = os.path.join(ROOT, "target/release/takana")
STORE = os.environ.get("HAMMER_STORE", os.path.join(ROOT, "store"))


def fatal(msg):
    print(f"✗ {msg}", file=sys.stderr)
    raise SystemExit(1)


def artefacto(receta, nombre, escotilla=None):
    # Misma escotilla que los guardianes: deja medir una copia instrumentada sin tocar el store.
    forzado = os.environ.get(escotilla) if escotilla else None
    if forzado:
        if not os.path.isdir(forzado):
            fatal(f"{escotilla}={forzado} no existe")
        print(f"⚠ {escotilla}: se mide {os.path.basename(forzado)}, que NO es el vigente")
        return forzado
    r = subprocess.run([HAMMER, "--store", STORE, "hash", f"recipes/{receta}.toml"],
                       cwd=ROOT, capture_output=True, text=True)
    if r.returncode != 0:
        fatal(f"no pude hashear {receta}: {r.stderr.strip()}")
    h = r.stdout.strip().splitlines()[-1].removeprefix("b3:")
    d = os.path.join(STORE, f"{h}-{nombre}")
    if not os.path.isdir(d) or not os.listdir(d):
        fatal(f"{nombre} b3:{h[:8]} no está sellado (o está vacío) — construilo antes")
    return d


def correr(capas, salida, url, segundos, capturar):
    """Un sway headless, un atuq dentro, y —si se pide— una foto antes de salir."""
    # La espera es la MISMA con o sin foto: en la corrida de descubrimiento también hay que darle
    # al navegador el tiempo de arrancar y hablar. La primera versión dormía 5 s y cerraba sway
    # antes de que la extensión dijera una palabra — el log salía vacío y parecía que no cargaba.
    foto = f"  sleep {segundos}\n"
    if capturar:
        foto += ("  /usr/bin/grim /salida/captura.png 2>>/salida/grim.log"
                 " || echo \"@@ grim falló\" >> /salida/grim.log\n")
    dentro = f"""#!/bin/sh
export XDG_RUNTIME_DIR=/run/user/0 HOME=/salida/hogar LANG=C XDG_DATA_DIRS=/usr/share
export WLR_BACKENDS=headless WLR_HEADLESS_OUTPUTS=1 WLR_RENDERER=pixman
mkdir -p "$XDG_RUNTIME_DIR" /salida/hogar; chmod 700 "$XDG_RUNTIME_DIR"
(
  i=0
  while [ ! -S /run/user/0/wayland-1 ] && [ $i -lt 40 ]; do sleep 1; i=$((i+1)); done
  [ -S /run/user/0/wayland-1 ] || {{ echo "@@ sin socket wayland" >> /salida/atuq.log; exit 1; }}
  export WAYLAND_DISPLAY=wayland-1
  /usr/bin/swaybg -c '#1a4b8c' & sleep 2
  export MOZ_ENABLE_WAYLAND=1 GDK_BACKEND=wayland LIBGL_DRIVERS_PATH=/usr/lib/dri
  export MOZ_DISABLE_CONTENT_SANDBOX=1 MOZ_DISABLE_RDD_SANDBOX=1 MOZ_DISABLE_UTILITY_SANDBOX=1
  /usr/bin/atuq --no-remote --profile /salida/perfil '{url}' >> /salida/atuq.log 2>&1 &
{foto}
  SWAYSOCK=$(ls /run/user/0/sway-ipc.* 2>/dev/null | head -1); export SWAYSOCK
  /usr/bin/swaymsg exit 2>/dev/null
) &
exec /usr/bin/sway -d > /salida/sway.log 2>&1
"""
    guion = os.path.join(salida, "dentro.sh")
    with open(guion, "w") as fh:
        fh.write(dentro)
    orden = ["bwrap", "--overlay-src", ALPINE, "--overlay-src", RFS]
    for c in capas:
        orden += ["--overlay-src", c]
    orden += [
        "--tmp-overlay", "/",
        "--dev-bind", "/dev", "/dev", "--proc", "/proc", "--bind", "/sys", "/sys",
        "--tmpfs", "/tmp", "--tmpfs", "/run", "--unshare-pid", "--unshare-net",
        "--bind", salida, "/salida",
        "--uid", "0", "--gid", "0",
        "--setenv", "PATH", "/usr/bin:/bin:/usr/sbin:/sbin",
        "sh", "/salida/dentro.sh",
    ]
    subprocess.run(orden, check=False, timeout=segundos + 120)
    log = os.path.join(salida, "atuq.log")
    return open(log, errors="replace").read() if os.path.isfile(log) else ""


def main():
    ap = argparse.ArgumentParser(description="captura de la barra lateral de IA")
    ap.add_argument("--question", help="pregunta para el modelo local (§6.7)")
    ap.add_argument("--archive", help="pregunta para tus páginas archivadas (§6.3)")
    ap.add_argument("--url", help="abrir esta URL en vez del panel (para aislar qué viene del arranque)")
    ap.add_argument("--out", default="work/atuq-ia.png", help="PNG de salida")
    ap.add_argument("--wait", type=int, default=75, help="segundos antes de la foto")
    args = ap.parse_args()
    if not args.question and not args.archive and not args.url:
        fatal("pasá --question, --archive o --url")

    if not os.path.isdir(RFS):
        fatal(f"no está el rootfs de sway en {RFS}")
    capas = [artefacto("atuq", "atuq", "ATUQ_DIR"), artefacto("puriy-costura", "puriy-costura"),
             artefacto("llama-cpp", "llama-cpp")]
    if not args.url:
        capas.append(artefacto("ia-modelo-embeddings" if args.archive else "ia-modelo-chat",
                               "ia-modelo-embeddings" if args.archive else "ia-modelo-chat"))
    for c in capas:
        print(f"   {os.path.basename(c)}")

    tmp = tempfile.mkdtemp(prefix="atuq-captura-")
    salida = os.path.join(tmp, "salida")
    os.makedirs(os.path.join(salida, "perfil"))
    try:
        with open(os.path.join(salida, "perfil", "user.js"), "w") as fh:
            fh.write('user_pref("browser.dom.window.dump.enabled", true);\n')
            fh.write('user_pref("browser.shell.checkDefaultBrowser", false);\n')

        # 1ª corrida: la extensión dice su URL base (Gecko la asigna por perfil).
        log = correr(capas, salida, "about:blank", 30, capturar=False)
        base = next((l[len("IA BASE "):].strip() for l in log.splitlines()
                     if l.startswith("IA BASE ")), None)
        if not base:
            destino = os.path.join(ROOT, "work", "atuq-captura.log")
            shutil.copy(os.path.join(salida, "atuq.log"), destino)
            fatal(f"la extensión no dijo su URL base — log en {destino}")
        print(f"\nextensión en {base}")

        # ⚠ El candado del perfil sobrevive a la primera corrida. Sway se lleva a `atuq` sin
        # despedirse, así que queda un `.parentlock` y la segunda corrida abre el diálogo
        # «Close Firefox» en vez de la página — lo dijo la primera captura, que salió con el
        # diálogo y la página en blanco. El perfil NO se puede tirar y rehacer: la URL base de la
        # extensión vive en él, y con un perfil nuevo el UUID cambia y la captura apuntaría a otro
        # lado.
        for candado in ("perfil/.parentlock", "perfil/lock", "perfil/parent.lock"):
            ruta = os.path.join(salida, candado)
            if os.path.lexists(ruta):
                os.remove(ruta)

        if args.url:
            url = args.url
        else:
            pregunta = args.question or args.archive
            clave = "q" if args.question else "p"
            url = f"{base}panel.html?{clave}=" + pregunta.replace(" ", "%20").replace("¿", "%C2%BF")
        log = correr(capas, salida, url, args.wait, capturar=True)
        for l in log.splitlines():
            if l.startswith("IA "):
                print("   " + l[:160])

        foto = os.path.join(salida, "captura.png")
        if not os.path.isfile(foto) or os.path.getsize(foto) == 0:
            g = os.path.join(salida, "grim.log")
            extra = open(g, errors="replace").read().strip() if os.path.isfile(g) else ""
            fatal(f"no salió la captura — {extra or 'grim no dejó rastro'}")
        destino = os.path.join(ROOT, args.out)
        os.makedirs(os.path.dirname(destino), exist_ok=True)
        shutil.copy(foto, destino)
        # El log viaja CON la foto, siempre y no sólo cuando algo falla: una captura muestra que
        # algo se ve raro y no por qué. La primera de estas capturas salió con tres barras de
        # notificación vacías, y sin el log al lado no había con qué empezar a mirar.
        log_destino = os.path.splitext(destino)[0] + ".log"
        shutil.copy(os.path.join(salida, "atuq.log"), log_destino)
        print(f"   log del navegador junto a la captura: {os.path.relpath(log_destino, ROOT)}")
        print(f"\n✓ captura en {args.out} ({os.path.getsize(destino)} bytes)")
        return 0
    finally:
        shutil.rmtree(tmp, ignore_errors=True)


sys.exit(main())