"""Neuronal su ZeroGPU: qui gira SOLO il modello.
La pagina, gli strumenti, la voce, i ricordi e le sessioni stanno nello Space col container
(amogaddy/neuronal). Quando quello deve calcolare chiama `sessione`: qui parte il modello sulla GPU,
si collega al container, fa i calcoli che gli vengono chiesti e dopo pochi secondi restituisce la GPU.

All'accensione, sulla CPU (non consuma tempo di GPU): scarica llama.cpp e il modello, compila neuronal-brain."""
import asyncio, glob, hmac, os, re, subprocess, tarfile, time, urllib.request

import gradio as gr
import spaces
from huggingface_hub import hf_hub_download

from ponte import ponte

QUI = os.path.dirname(os.path.abspath(__file__))
BASE = "/tmp/neuronal"
LLAMA = "https://github.com/ggml-org/llama.cpp/releases/download/b11200/"
PACCHETTI = ["llama-b11200-bin-ubuntu-cuda-12.8-x64.tar.gz", "cudart-llama-b11200-bin-ubuntu-cuda-12.8-x64.tar.gz"]
# llama.cpp è compilato per una libreria di sistema più nuova di quella dello Space: se ne affianca una copia
UBUNTU = "http://archive.ubuntu.com/ubuntu/pool/main/g/"
DEB = [("glibc/", r"libc6_2\.39-0ubuntu8[.\d]*_amd64\.deb"), ("gcc-14/", r"libstdc\+\+6_14[^\"]*24\.04[.\d]*_amd64\.deb"),
       ("gcc-14/", r"libgcc-s1_14[^\"]*24\.04[.\d]*_amd64\.deb"), ("gcc-14/", r"libgomp1_14[^\"]*24\.04[.\d]*_amd64\.deb")]
CHIAVE = os.environ.get("NEURONAL_CHIAVE", "")
MASSIMO = 110                                # durata massima di una sessione di GPU, in secondi
REGISTRO = os.path.join(BASE, "sessioni.log")


def prepara():
    os.makedirs(BASE, exist_ok=True)
    t = time.time()
    for nome in PACCHETTI:
        fatto = os.path.join(BASE, nome + ".ok")
        if not os.path.exists(fatto):
            arch = os.path.join(BASE, nome)
            urllib.request.urlretrieve(LLAMA + nome, arch)
            with tarfile.open(arch) as tar:
                tar.extractall(os.path.join(BASE, "llama"))
            os.remove(arch)
            open(fatto, "w").close()
    libs = sorted({d for d, _, fs in os.walk(os.path.join(BASE, "llama")) if any(".so" in f for f in fs)})
    llama = next(d for d in libs if os.path.exists(os.path.join(d, "libllama.so")))

    sistema = os.path.join(BASE, "glibc")
    if not os.path.exists(os.path.join(sistema, ".ok")):
        for cartella, modello in DEB:
            elenco = urllib.request.urlopen(UBUNTU + cartella, timeout=60).read().decode()
            nome = sorted(set(re.findall(modello, elenco)))[-1]
            deb = os.path.join(BASE, nome)
            urllib.request.urlretrieve(UBUNTU + cartella + nome, deb)
            subprocess.run(["dpkg-deb", "-x", deb, sistema], check=True)
            os.remove(deb)
        open(os.path.join(sistema, ".ok"), "w").close()
    sis = os.path.join(sistema, "usr", "lib", "x86_64-linux-gnu")
    if not os.path.exists(os.path.join(sis, "libc.so.6")):
        sis = os.path.join(sistema, "lib", "x86_64-linux-gnu")

    cervello = os.path.join(BASE, "neuronal-brain")
    src = os.path.join(QUI, "src")
    subprocess.run(["gcc", "-O2", "-std=gnu11", "-o", cervello, os.path.join(src, "brain_worker.c"), os.path.join(src, "brain.c"),
                    "-I" + src, "-I" + os.path.join(src, "llama"), "-L" + llama, "-lmtmd", "-lllama", "-lggml", "-lggml-base",
                    "-lm", "-lpthread", "-Wl,--allow-shlib-undefined"], check=True)

    modello = hf_hub_download("prism-ml/Bonsai-27B-gguf", "Bonsai-27B-Q1_0.gguf")
    visione = hf_hub_download("prism-ml/Bonsai-27B-gguf", "Bonsai-27B-mmproj-Q8_0.gguf")
    print(f"Pronto in {time.time() - t:.0f} s: llama.cpp, modello e neuronal-brain.", flush=True)
    percorsi = os.pathsep.join([sis] + libs + ["/usr/lib/x86_64-linux-gnu"])
    return [os.path.join(sis, "ld-linux-x86-64.so.2"), "--library-path", percorsi, cervello, modello, llama, visione, "32768"]


COMANDO = prepara()


def valida(chiave):
    return bool(CHIAVE) and hmac.compare_digest(str(chiave), CHIAVE)


def limita(secondi):
    try:
        return max(10, min(int(secondi), MASSIMO))
    except (TypeError, ValueError):
        return 10


def durata(chiave, url, secondi):
    """Quanta GPU chiedere: il tempo della sessione più il margine per chiudere. ZeroGPU concede la sessione
    solo se nel tempo del giorno ne resta almeno altrettanto, ma conta solo quello usato davvero."""
    return limita(secondi) + 8


@spaces.GPU(duration=durata)
def sessione(chiave, url, secondi):
    """Avvia il modello sulla GPU e lo collega al motore di Neuronal che l'ha chiesto."""
    if not valida(chiave):
        return "chiave sbagliata"
    if not re.fullmatch(r"wss?://[A-Za-z0-9.\-:]+/ws\?role=brain&ticket=[0-9a-f]{16,40}", str(url)):
        return "indirizzo non valido"
    secondi = limita(secondi)
    with open(REGISTRO, "ab") as reg:
        reg.write(f"\n--- {time.strftime('%H:%M:%S')} sessione di {secondi} s\n".encode())
        reg.flush()
        try:
            usati = asyncio.run(ponte(COMANDO, url, secondi, stderr=reg))
        except Exception as e:
            reg.write(f"ponte: {type(e).__name__}: {e}\n".encode())
            return f"collegamento al motore non riuscito ({type(e).__name__})"
    return f"ok {usati:.1f}"


def diagnosi(chiave):
    """Le ultime righe scritte dal modello (per capire un problema senza usare la GPU)."""
    if not valida(chiave):
        return "chiave sbagliata"
    try:
        return open(REGISTRO, encoding="utf-8", errors="replace").read()[-6000:]
    except OSError:
        return "nessuna sessione finora"


with gr.Blocks(title="Neuronal GPU") as pagina:
    gr.Markdown("# 🧠 Neuronal — il modello\nQui gira solo il modello di Neuronal, sulla GPU. "
                "Per usarlo vai su **[amogaddy/neuronal](https://huggingface.co/spaces/amogaddy/neuronal)**: "
                "la pagina, gli strumenti e la voce sono lì.\n\nMade in Italy 🇮🇹")
    with gr.Accordion("Servizio (lo usa il motore di Neuronal)", open=False):
        chiave = gr.Textbox(label="Chiave", type="password")
        url = gr.Textbox(label="Motore")
        secondi = gr.Number(label="Secondi", value=50)
        esito = gr.Textbox(label="Esito", lines=12)
        gr.Button("Sessione").click(sessione, [chiave, url, secondi], esito, api_name="sessione")
        gr.Button("Diagnosi").click(diagnosi, [chiave], esito, api_name="diagnosi")

pagina.launch(show_error=True)