#!/usr/bin/env python3
"""¿La barra lateral le pregunta a un modelo que corre EN ESTA MÁQUINA? (SDD 26 §6.7)

    python3 scripts/test-atuq-ia.py                      # sale una respuesta generada acá
    python3 scripts/test-atuq-ia.py --negative-control    # sin modelo: lo DICE, no inventa
    python3 scripts/test-atuq-ia.py --image-model         # con el modelo DE LA IMAGEN (1 GiB, lento)

── QUÉ SE EJERCITA, DE PUNTA A PUNTA ─────────────────────────────────────────────────────────
El navegador abre el panel de la extensión con una pregunta (`panel.html?q=…`), el panel se la pasa
al fondo, el fondo al host por native messaging, el host levanta un `llama-server` —que NO estaba
corriendo: no hay servicio de IA en la imagen— y la respuesta vuelve por el mismo camino. La
inferencia es REAL: tokens generados por `llama-cpp` del corpus, contra un modelo pineado por
sha256, dentro de la jaula y sin red.

El modelo de prueba es el MISMO que pinea `scripts/test-llama-cpp.py`, importado de ahí para que el
sha256 viva en un solo sitio: `stories260K`, 1,1 MB, que dice tonterías con gramática — que es
exactamente lo que hace falta para comprobar el CAMINO sin pinear un modelo de producción, que es
otra decisión (y otros cientos de megas).

── EL CONTROL NEGATIVO ───────────────────────────────────────────────────────────────────────
`--negative-control` borra el modelo de la jaula y exige que el panel muestre **la causa**. Una IA
que ante un fallo contesta algo genérico es peor que una que no contesta: el usuario no puede
distinguir «no hay modelo» de «el modelo dijo eso».

REQUISITOS: el rootfs mínimo de `atuq-nested.sh`, y `atuq` + `puriy-costura` + `llama-cpp` sellados.
"""
import importlib.util
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/atuq-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"))
# El path de PRODUCCIÓN del modelo, el mismo que la constante del host. Se mide éste y no la
# escotilla de pruebas: una escotilla mide el código, no el contrato con la imagen.
MODELO_EN_IMAGEN = "/usr/share/takana/ia/modelo.gguf"
PREGUNTA = "Once upon a time"


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


def artefacto(receta, nombre, escotilla=None):
    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 necesariamente 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):
        fatal(f"{nombre} b3:{h[:8]} no está sellado — construilo antes")
    if not os.listdir(d):
        fatal(f"{nombre} b3:{h[:8]} está VACÍO: eso no es un artefacto, es un nombre")
    return d


def modelo_de_prueba():
    """El fetcher del guardián del motor, reusado: así el sha256 del modelo vive en UN sitio."""
    ruta = os.path.join(ROOT, "scripts", "test-llama-cpp.py")
    spec = importlib.util.spec_from_file_location("guardian_llama", ruta)
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)      # tiene guarda `__main__`: importarlo no corre nada
    return str(mod.bajar_modelo())


def correr(atuq, costura, llama, capa_modelo, salida, url, negativo, tope=90):
    borrar = f"rm -f {MODELO_EN_IMAGEN}" if negativo else "true"
    # El modelo de la imagen llega como CAPA (su artefacto ya lo pone en el path de producción); el
    # de juguete se copia, que son 1,1 MB.
    poner_modelo = "true" if capa_modelo else f"cp /salida/modelo.gguf {MODELO_EN_IMAGEN}"
    guion = f"""
    mkdir -p /root /salida/hogar /salida/run {os.path.dirname(MODELO_EN_IMAGEN)}
    [ -s /etc/machine-id ] || tr -d - < /proc/sys/kernel/random/uuid > /etc/machine-id
    export HOME=/salida/hogar XDG_RUNTIME_DIR=/salida/run
    {poner_modelo}
    {borrar}
    timeout {tope} /usr/bin/atuq --profile /salida/perfil --no-remote --headless '{url}' \
        > /salida/consola.log 2>&1
    # ¿Se fue el motor con el navegador? Se cuenta ACÁ ADENTRO, con el namespace de PIDs todavía
    # vivo: desde el hub no se vería nada —`--unshare-pid` se lleva todo al cerrar la jaula— y esa
    # limpieza del arnés taparía justo la promesa que hay que medir. Medio giga de modelo que
    # sobrevive al navegador no se nota hasta que la máquina se queda sin RAM.
    # Se cuenta leyendo /proc a mano, y no con `pgrep -c`: busybox NO tiene esa opción, y el
    # `|| echo 0` que parecía prudente convertía el error en un CERO — o sea en un veredicto verde.
    # Medido: con un motor vivo a propósito, el guardián pasaba igual.
    # Se espera a que SE MUERA, hasta 15 s, y se anota cuánto tardó. Matar a un proceso con un
    # giga de modelo mapeado no es instantáneo, y contar en el instante exacto del cierre confundía
    # «tardando en morir» con «fuga»: con el modelo de juguete daba 0 y con el de la imagen daba 1.
    esperado=0
    while [ $esperado -lt 15 ]; do
        n=0
        for p in /proc/[0-9]*; do
            [ "$(cat $p/comm 2>/dev/null)" = "llama-server" ] && n=$((n+1))
        done
        [ "$n" = "0" ] && break
        esperado=$((esperado+1))
        sleep 1
    done
    echo "$n $esperado" > /salida/motores-vivos
    exit 0
    """
    subprocess.run([
        "bwrap",
        "--overlay-src", ALPINE, "--overlay-src", RFS,
        "--overlay-src", atuq, "--overlay-src", costura, "--overlay-src", llama,
        *(["--overlay-src", capa_modelo] if capa_modelo else []),
        "--tmp-overlay", "/",
        "--dev-bind", "/dev", "/dev", "--proc", "/proc", "--bind", "/sys", "/sys",
        "--tmpfs", "/tmp", "--unshare-pid", "--unshare-net",
        "--bind", salida, "/salida",
        "--uid", "0", "--gid", "0",
        "--setenv", "PATH", "/usr/bin:/bin:/usr/sbin:/sbin",
        "--setenv", "XDG_DATA_DIRS", "/usr/share", "--setenv", "MOZ_HEADLESS", "1",
        "--setenv", "MOZ_DISABLE_CONTENT_SANDBOX", "1",
        "--setenv", "MOZ_DISABLE_GMP_SANDBOX", "1",
        "--setenv", "MOZ_DISABLE_RDD_SANDBOX", "1",
        "--setenv", "MOZ_DISABLE_SOCKET_PROCESS_SANDBOX", "1",
        "--setenv", "MOZ_DISABLE_UTILITY_SANDBOX", "1",
        "sh", "-c", guion,
    ], check=False)
    log = open(os.path.join(salida, "consola.log"), errors="replace").read()
    vivos = "?"
    try:
        vivos = open(os.path.join(salida, "motores-vivos"), errors="replace").read().strip()
    except OSError:
        pass
    return [l[len("IA "):] for l in log.splitlines() if l.startswith("IA ")], log, vivos


def main():
    negativo = "--negative-control" in sys.argv[1:]
    # Por defecto se mide con el modelo de juguete: 1,1 MB, instantáneo, y alcanza para certificar la
    # CADENA. Con `--image-model` se mide el de la imagen — la misma cadena, pero contestando de
    # verdad; cuesta 1 GiB de overlay y unos segundos de generación, y por eso no es el default.
    de_la_imagen = "--image-model" in sys.argv[1:]
    if negativo:
        print("== MODO CONTROL NEGATIVO: sin modelo en la imagen; el panel tiene que DECIR la causa")
    if not os.path.isdir(RFS):
        fatal(f"no está el rootfs en {RFS} — corré antes `scripts/atuq-nested.sh`")
    atuq = artefacto("atuq", "atuq", "ATUQ_DIR")
    costura = artefacto("puriy-costura", "puriy-costura")
    llama = artefacto("llama-cpp", "llama-cpp", "LLAMA_DIR")
    print(f"artefactos: {os.path.basename(atuq)}\n            {os.path.basename(costura)}"
          f"\n            {os.path.basename(llama)}")
    capa_modelo = None
    modelo = None
    if de_la_imagen:
        # Llega como CAPA: su artefacto ya instala el modelo en el path de producción. Copiarlo
        # sería mover 1 GiB a un tmpfs — lento, y la primera vez llenó /tmp y el guardián murió con
        # un ENOSPC que no tenía nada que ver con lo que mide.
        capa_modelo = artefacto("ia-modelo-chat", "ia-modelo-chat")
        print(f"modelo: el de la imagen — {os.path.basename(capa_modelo)}")
    else:
        modelo = modelo_de_prueba()
        print(f"modelo de prueba: {os.path.basename(modelo)}")

    tmp = tempfile.mkdtemp(prefix="atuq-ia-")
    salida = os.path.join(tmp, "salida")
    os.makedirs(os.path.join(salida, "perfil"))
    if modelo:
        shutil.copy(modelo, os.path.join(salida, "modelo.gguf"))
    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')

        # Primera corrida: descubrir la URL base que Gecko le asignó a la extensión (es por perfil).
        lineas, log, _ = correr(atuq, costura, llama, capa_modelo, salida, "about:blank", negativo)
        base = next((l[len("BASE "):].strip() for l in lineas if l.startswith("BASE ")), None)
        if not base:
            destino = os.path.join(ROOT, "work", "atuq-ia.log")
            os.makedirs(os.path.dirname(destino), exist_ok=True)
            shutil.copy(os.path.join(salida, "consola.log"), destino)
            fatal(f"la extensión no dijo su URL base — log en {destino}")
        print(f"\nextensión en {base}")

        panel = f"{base}panel.html?q={PREGUNTA.replace(' ', '%20')}"
        # 512 tokens a ~20 tok/s son ~25 s, más la carga del modelo: el tope se acompaña solo.
        lineas, log, motores = correr(atuq, costura, llama, capa_modelo, salida, panel, negativo,
                                      tope=180 if de_la_imagen else 90)
        for l in lineas:
            print("   " + l)

        errores = [l for l in lineas if l.startswith("ERROR ")]
        if negativo:
            if not errores:
                fatal("sin modelo no dijo nada: el panel se habría quedado pensando para siempre")
            if "no hay modelo" not in errores[0]:
                fatal(f"la causa no nombra lo que falta: {errores[0]}")
            if any(l.startswith("RESPUESTA ") for l in lineas):
                fatal("sin modelo CONTESTÓ igual: eso es inventar una respuesta")
            print("\n✓ control negativo: sin modelo lo dice, nombrando lo que falta")
            print("✓ y no inventa una respuesta, que es el único fallo inaceptable en esto")
            return 0

        if errores:
            fatal(f"el host falló: {errores[0]}")
        respuestas = [l for l in lineas if l.startswith("RESPUESTA ")]
        if not respuestas:
            fatal("no volvió ninguna respuesta del modelo")
        r = respuestas[0]
        if f"modelo={MODELO_EN_IMAGEN}" not in r:
            fatal(f"la respuesta no sale del modelo de la imagen: {r}")
        # Un `texto=""` es lo que devolvería un camino que «funciona» sin generar nada.
        texto = r.split("texto=", 1)[1] if "texto=" in r else ""
        if len(texto.strip(' "')) < 2:
            fatal(f"la respuesta vino VACÍA: {r}")
        quedaron, esperado = (motores.split() + ["?"])[:2]
        if quedaron != "0":
            fatal(f"quedaron {quedaron} `llama-server` vivos {esperado}s después de cerrarse el"
                  " navegador — el motor tiene que irse con él (§6.7): medio giga de modelo huérfano"
                  " no se ve hasta que la máquina se queda sin RAM")
        print("\n✓ el navegador preguntó y un modelo de ESTA máquina contestó")
        print("✓ y nadie levantó el servidor: no hay servicio de IA en la imagen — lo levantó el host")
        print("✓ la jaula corre con `--unshare-net`: no hubo red para consultar a nadie más")
        print(f"✓ y el motor se fue con el navegador: 0 `llama-server` vivos ({esperado}s después)")
        return 0
    finally:
        shutil.rmtree(tmp, ignore_errors=True)


sys.exit(main())