#!/usr/bin/env python3
"""
colibrì — piccolo motore, modello immenso.
CLI per far girare GLM-5.2 (744B) in locale, su CPU, in ~15-26 GB di RAM.

  coli chat                 chat interattiva (carica il modello UNA volta)
  coli run "prompt"         generazione singola
  coli info                 stato: modello, RAM, disco, config
  coli bench [task...]      benchmark di qualità (MMLU/HellaSwag/...)
  coli convert              converte GLM-5.2-FP8 -> int4 (streaming)
  coli models [add <dir>]   registro multi-modello (scelta all'avvio della chat)
  coli build                compila il motore

Config via env o flag (validi anche dopo il sottocomando):
  COLI_MODEL=<dir>   modello (default /home/vincenzo/glm52_i4)
  --ram N            budget RAM in GB (auto-cap cache expert)
  --topp P           top-p adattivo sugli expert     --topk N   top-k fisso
  --ngen N           token massimi per risposta      --cap N    slot cache/layer
"""
import os, sys, subprocess, argparse, json, time, signal, shutil, threading, re, codecs, tempfile, textwrap

HERE = os.path.dirname(os.path.abspath(__file__))
GLM  = os.path.join(HERE, "glm")
DEF_MODEL = os.environ.get("COLI_MODEL", "/home/vincenzo/glm52_i4")
END   = b"\x01\x01END\x01\x01\n"
READY = b"\x01\x01READY\x01\x01\n"

# ---------- palette & stile ----------
def _c(n): return f"\033[38;5;{n}m"
class C:
    teal=_c(37); cyan=_c(80); mag=_c(170); org=_c(208); grn=_c(78); yel=_c(179)
    dim="\033[2m"; b="\033[1m"; r="\033[0m"; gray=_c(242); dgray=_c(238)
    @staticmethod
    def off():
        for k,v in vars(C).items():
            if isinstance(v,str) and v.startswith("\033"): setattr(C,k,"")
TTY = sys.stdout.isatty() or os.environ.get("COLI_COLOR")=="1"
if not TTY: C.off()

# ---------- colibrì 8-bit (pixel art, 2 pixel verticali per carattere) ----------
SPRITE = [
    "....MMM.........",
    "...MMMMM..w.....",
    "....MMMM.ww.....",
    "OOOOTTeTCC......",
    "....TTTTTCC.....",
    ".....TTTTCC.....",
    "......TTCC......",
    ".......TC.......",
    "........C.......",
    "................",
]
PAL = {"M":170, "T":37, "C":80, "O":208, "e":231, "w":80, ".":None}

def sprite_lines():
    if not TTY:
        return ["  (\\   ", "   )·>  ", "  / \\   ", "        ", "        "]
    out=[]
    for y in range(0,len(SPRITE),2):
        top, bot = SPRITE[y], SPRITE[y+1] if y+1<len(SPRITE) else "."*len(SPRITE[y])
        row=""
        for x in range(len(top)):
            ct, cb = PAL.get(top[x]), PAL.get(bot[x])
            if ct is None and cb is None: row+= "\033[0m "
            elif ct is not None and cb is None: row+= f"\033[38;5;{ct}m\033[49m▀"
            elif ct is None and cb is not None: row+= f"\033[38;5;{cb}m\033[49m▄"
            else: row+= f"\033[38;5;{ct}m\033[48;5;{cb}m▀"
        out.append(row+"\033[0m")
    return out

def banner(sub=""):
    sp=sprite_lines()
    txt=[
        f"{C.teal}{C.b}colibrì{C.r} {C.dim}v1.0{C.r}",
        f"{C.dim}piccolo motore, modello immenso{C.r}",
        f"{C.gray}GLM-5.2 · 744B MoE · int4 · streaming CPU{C.r}",
        f"{C.dgray}{sub}{C.r}" if sub else "",
        "",
    ]
    print()
    for i,s in enumerate(sp):
        t = txt[i] if i<len(txt) else ""
        print(f"  {s}   {t}")
    print(f"  {C.dgray}{'─'*58}{C.r}")

def hline(w): return f"{C.dgray}{'─'*w}{C.r}"

# ---------- util ----------
def term_w(): return min(shutil.get_terminal_size((80,20)).columns, 100)

# ---------- registro multi-modello ----------
# Ogni ARCHITETTURA ha il SUO motore (file C e binario separati): i modelli non si
# toccano tra loro. Il registro e' models.json accanto a coli: {nome: {dir, arch}}.
#   coli models                     lista
#   coli models add <dir> [nome]    registra (arch letta da config.json)
#   coli chat                       1 modello -> parte; piu' modelli -> scegli all'avvio
ENGINES = {                          # arch (config.json) -> binario del motore
    "glm_moe_dsa": "glm",            # GLM-5.2 (MLA + DSA + MTP)          — PRONTO
    "gpt_oss":     "gptoss",         # gpt-oss-120b (GQA+sinks, MoE 128e) — in arrivo
    "qwen3_5_moe": "qwenmoe",        # Qwen3.6-A3B (MoE 256e)             — in arrivo
}
def model_arch(d):
    try:
        c=json.load(open(os.path.join(d,"config.json")))
        return c.get("model_type") or c.get("text_config",{}).get("model_type") or "?"
    except Exception: return "?"
def models_load():
    reg={}
    p=os.path.join(HERE,"models.json")
    if os.path.exists(p):
        try: reg=json.load(open(p))
        except Exception: reg={}
    # il modello di default e' sempre registrato se esiste su disco
    if DEF_MODEL and os.path.isdir(DEF_MODEL) and not any(v.get("dir")==DEF_MODEL for v in reg.values()):
        reg.setdefault(os.path.basename(DEF_MODEL), {"dir":DEF_MODEL, "arch":model_arch(DEF_MODEL)})
    return reg
def models_save(reg):
    json.dump(reg, open(os.path.join(HERE,"models.json"),"w"), indent=1)
def engine_for(model_dir):
    arch=model_arch(model_dir)
    binname=ENGINES.get(arch)
    if not binname:
        sys.exit(f"{C.yel}architettura '{arch}' non (ancora) supportata.{C.r} Motori pronti: "
                 +", ".join(f"{k}->{v}" for k,v in ENGINES.items()))
    path=os.path.join(HERE,binname)
    if not os.path.exists(path):
        sys.exit(f"{C.yel}motore '{binname}' non compilato per l'architettura {arch}.{C.r}"
                 +("  Esegui: coli build" if binname=="glm" else "  (motore in sviluppo)"))
    return path
def pick_model(a):
    """se --model/COLI_MODEL e' esplicito usa quello; altrimenti: 1 registrato -> quello,
    piu' d'uno -> menu di scelta all'avvio della chat (zero passaggi extra)."""
    if a.model != DEF_MODEL or not sys.stdin.isatty():
        return a.model
    reg=models_load()
    avail=[(n,v) for n,v in reg.items() if os.path.isdir(v.get("dir",""))]
    if len(avail)<=1:
        return avail[0][1]["dir"] if avail else a.model
    print(f"  {C.dim}modelli disponibili:{C.r}")
    for i,(n,v) in enumerate(avail):
        eng=ENGINES.get(v.get("arch","?"),"?")
        ok = "✓" if os.path.exists(os.path.join(HERE,eng)) else "motore in sviluppo"
        print(f"    {C.teal}{i+1}{C.r}) {n}  {C.dgray}· {v.get('arch','?')} · {ok}{C.r}")
    try: ch=input(f"  {C.teal}›{C.r} quale? [1] ").strip() or "1"
    except EOFError: ch="1"
    try: return avail[max(0,min(len(avail)-1,int(ch)-1))][1]["dir"]
    except ValueError: return avail[0][1]["dir"]

def need_model(model):
    if not os.path.isdir(model):
        sys.exit(f"{C.yel}modello non trovato:{C.r} {model}\n  imposta COLI_MODEL o usa --model")
    if not os.path.exists(os.path.join(model,"tokenizer.json")):
        sys.exit(f"{C.yel}manca tokenizer.json in {model}{C.r}")
    engine_for(model)                 # verifica che il motore per QUESTA architettura esista

def env_for(a):
    e = dict(os.environ, SNAP=a.model)
    if a.ram:  e["RAM_GB"]=str(a.ram)
    if a.ngen: e["NGEN"]=str(a.ngen)
    if a.topp: e["TOPP"]=str(a.topp)
    if a.topk: e["TOPK"]=str(a.topk)
    if a.temp is not None: e["TEMP"]=str(a.temp)   # 0 = greedy; default motore: 1.0 + nucleus 0.95
    return e

class Spinner:
    FRAMES=["⠋","⠙","⠹","⠸","⠼","⠴","⠦","⠧","⠇","⠏"]
    def __init__(self,label,tick=None):
        self.label=label; self.tick=tick; self.suffix=""
        self.stop_evt=threading.Event(); self.t0=time.time(); self.th=None
    def start(self):
        if not TTY: return
        def run():
            i=0
            while not self.stop_evt.is_set():
                el=time.time()-self.t0
                if self.tick and i%8==0:              # ~1 Hz: legge il progresso dal log
                    try: self.suffix=self.tick() or self.suffix
                    except Exception: pass
                suf=f" {C.dgray}· {self.suffix}{C.r}" if self.suffix else ""
                sys.stdout.write(f"\r  {C.teal}{self.FRAMES[i%10]}{C.r} {C.dim}{self.label} {el:.0f}s{C.r}{suf}\033[K")
                sys.stdout.flush(); i+=1; time.sleep(0.12)
        self.th=threading.Thread(target=run,daemon=True); self.th.start()
    def stop(self):
        self.stop_evt.set()
        if self.th: self.th.join(timeout=0.4)
        if TTY: sys.stdout.write("\r\033[K"); sys.stdout.flush()

def stream_turn(p, sentinel, on_bytes):
    """legge fino alla sentinella; on_bytes riceve i chunk della risposta. Poi legge la riga STAT."""
    pend=b""
    while True:
        b=p.stdout.read(1)
        if b==b"": return None
        pend+=b
        if pend.endswith(sentinel):
            rest=pend[:-len(sentinel)]
            if rest: on_bytes(rest)
            line=p.stdout.readline().decode("utf-8","replace").strip()   # STAT tok tps hit rss
            m=re.match(r"STAT (\S+) (\S+) (\S+) (\S+)", line)
            return {"tok":int(m.group(1)),"tps":float(m.group(2)),"hit":float(m.group(3)),"rss":float(m.group(4))} if m else {}
        if len(pend)>len(sentinel):
            out=pend[:-len(sentinel)]; pend=pend[-len(sentinel):]
            on_bytes(out)

# ---------- comandi ----------
def cmd_build(a):
    banner("build")
    sys.exit(subprocess.call(["make","-C",HERE,"glm"]))

def cmd_info(a):
    banner("info")
    cfgp=os.path.join(a.model,"config.json")
    def row(k,v): print(f"   {C.gray}{k:<10}{C.r} {v}")
    if os.path.exists(cfgp):
        c=json.load(open(cfgp))
        row("modello", a.model)
        row("arch", f"hidden {c.get('hidden_size')} · {c.get('num_hidden_layers')} layer · "
                    f"{c.get('n_routed_experts')} expert/layer · top-{c.get('num_experts_per_tok')}")
        sts=[x for x in os.listdir(a.model) if x.endswith('.safetensors')]
        sz=sum(os.path.getsize(os.path.join(a.model,x)) for x in sts)
        row("shard", f"{len(sts)} file · {sz/1e9:.0f} GB su disco")
    else:
        print(f"   {C.yel}config.json non presente (conversione incompleta?){C.r}")
    try:
        mi=open('/proc/meminfo').read()
        tot=int(re.search(r'MemTotal:\s+(\d+)',mi).group(1))/1e6
        av=int(re.search(r'MemAvailable:\s+(\d+)',mi).group(1))/1e6
        row("RAM", f"{tot:.0f} GB totali · {av:.1f} GB disponibili")
    except Exception: pass
    fs=os.statvfs(a.model if os.path.isdir(a.model) else HERE)
    row("disco", f"{fs.f_bavail*fs.f_frsize/1e9:.0f} GB liberi")
    row("motore", "pronto ✓" if os.path.exists(GLM) else "da compilare (coli build)")
    knobs=[]
    if a.ram: knobs.append(f"ram {a.ram}GB")
    if a.topp: knobs.append(f"topp {a.topp}")
    if a.topk: knobs.append(f"topk {a.topk}")
    if knobs: row("tuning", " · ".join(knobs))
    print()

def cmd_run(a):
    a.model=pick_model(a)
    need_model(a.model)
    prompt=" ".join(a.prompt) if a.prompt else sys.exit('uso: coli run "il tuo prompt"')
    banner("run")
    # template ufficiale GLM-5.2: niente \n dopo i ruoli; <think></think> = risposta diretta (nothink)
    e=env_for(a); e["PROMPT"]=f"[gMASK]<sop><|user|>{prompt}<|assistant|><think></think>"
    sys.exit(subprocess.call([engine_for(a.model), str(a.cap)], env=e))

def cmd_chat(a):
    a.model=pick_model(a)
    need_model(a.model)
    banner(f"chat · {os.path.basename(a.model)} · ram {a.ram or '-'}GB · topp {a.topp or 'off'}")
    errlog=tempfile.NamedTemporaryFile(mode="w+", suffix=".log", delete=False)
    e=env_for(a); e["SERVE"]="1"
    p=subprocess.Popen([engine_for(a.model),str(a.cap)], env=e, stdin=subprocess.PIPE,
                       stdout=subprocess.PIPE, stderr=errlog, bufsize=0)
    sp=Spinner("sveglio il gigante (744B)…"); sp.start()
    st=stream_turn(p, READY, lambda b: None)
    sp.stop()
    if st is None:
        errlog.seek(0); print(errlog.read()[-1500:]); sys.exit("il motore è uscito durante il load")
    errlog.flush()
    try:
        elog=open(errlog.name).read()
        mload=re.search(r"caricato in ([0-9.]+)s \| densa residente: ([0-9.]+) MB", elog)
        if mload: print(f"  {C.grn}✓{C.r} pronto in {mload.group(1)}s {C.dim}· residente {float(mload.group(2))/1000:.1f} GB · RSS {st.get('rss','?')} GB{C.r}")
        for l in elog.splitlines():                     # una riga di stato per riga, senza path
            if l.startswith(("[RAM_GB","[PIN]","[MTP]","[USAGE]","[DSA]")):
                l=re.sub(r" ?\(?/[^ )]+\)?","",l.strip())       # via i percorsi lunghi
                l=re.sub(r" da$| in$","",l)
                for chunk in textwrap.wrap(l, term_w()-4) or [l]:
                    print(f"  {C.dgray}{chunk}{C.r}")
    except Exception: pass
    print(f"  {C.dim}scrivi e premi invio · :piu continua risposta · :reset memoria · :q esci{C.r}\n")
    w=term_w()-4
    def user_box(msg):
        """ri-disegna il messaggio dentro una box che si ADATTA su piu' righe:
        l'input grezzo (che sborda) viene cancellato e sostituito dal testo avvolto."""
        cols=shutil.get_terminal_size((80,20)).columns
        used=max(1, (6+len(msg)+cols-1)//cols)          # righe occupate dall'input ("  │ › "+msg)
        sys.stdout.write(f"\x1b[{used}A\x1b[0J")        # su di N righe e pulisci fino in fondo
        inner=w-3                                        # spazio utile: "  │ › "+testo+"│" = w+4 colonne
        lines=textwrap.wrap(msg, inner) or [""]
        for i,ln in enumerate(lines):
            pre = f"{C.teal}{C.b}›{C.r}" if i==0 else " "
            print(f"  {C.dgray}│{C.r} {pre} {ln}{' '*(inner-len(ln))}{C.dgray}│{C.r}")
        print(f"  {C.dgray}╰{'─'*w}╯{C.r}")
    try:
        while True:
            if TTY:
                print(f"  {C.dgray}╭{'─'*w}╮{C.r}")
                try: msg=input(f"  {C.dgray}│{C.r} {C.teal}{C.b}›{C.r} ")
                except EOFError: print(); break
                try: user_box(msg.strip())
                except Exception: print(f"  {C.dgray}╰{'─'*w}╯{C.r}")
            else:
                try: msg=input()
                except EOFError: break
            msg=msg.strip()
            if msg in (":q",":quit","exit"): break
            if not msg: continue
            if msg==":reset":
                p.stdin.write(b"\x02RESET\n"); p.stdin.flush()
                stream_turn(p, END, lambda b: None)
                print(f"  {C.dim}✦ memoria azzerata{C.r}\n"); continue
            if msg in (":piu",":più",":more",":continua"):
                p.stdin.write(b"\x02MORE\n"); p.stdin.flush()
            else:
                p.stdin.write((msg.replace("\n"," ")+"\n").encode()); p.stdin.flush()
            print(f"\n  {C.teal}◆ colibrì{C.r}")
            dec=codecs.getincrementaldecoder("utf-8")("replace")
            state={"first":True}
            def prefill_tick(path=errlog.name):
                try:
                    with open(path) as f:
                        f.seek(max(0, os.path.getsize(path)-1500)); tail=f.read()
                    pl=[l for l in tail.splitlines() if l.startswith("[prefill]")]
                    return pl[-1].replace("[prefill] ","prefill ") if pl else ""
                except Exception: return ""
            sp2=Spinner("pensa…", tick=prefill_tick); sp2.start()
            def echo(bs, _dec=dec, _st=state):
                if _st["first"]:
                    sp2.stop(); _st["first"]=False
                    sys.stdout.write("  ")
                s=_dec.decode(bs)
                if s: sys.stdout.write(s.replace("\n","\n  ")); sys.stdout.flush()
            t0=time.time()
            st=stream_turn(p, END, echo)
            sp2.stop()
            if st is None: print(f"\n  {C.yel}[motore terminato]{C.r}"); break
            el=time.time()-t0
            if st.get("tok"):
                print(f"\r  {C.dgray}└─ {st['tok']} tok · {st['tps']:.2f} tok/s · hit {st['hit']:.0f}% · RSS {st['rss']:.1f} GB · {el:.0f}s{C.r}")
                if st["tok"]>=a.ngen:
                    print(f"  {C.yel}…troncato al limite --ngen ({a.ngen}): scrivi :piu per far continuare la risposta{C.r}")
                print()
            else:
                print()
    except KeyboardInterrupt:
        print(f"\n  {C.dim}interrotto{C.r}")
    finally:
        try: p.stdin.close(); p.terminate()
        except Exception: pass
        try: os.unlink(errlog.name)
        except Exception: pass
    print(f"  {C.teal}ciao{C.r} {C.dim}— il colibrì torna al nido{C.r} 🐦\n")

def cmd_bench(a):
    need_model(a.model)
    banner("bench")
    # python con `tokenizers`: l'ambiente del progetto se c'e', altrimenti quello corrente
    venv_py=os.path.join(HERE,"mio_env","bin","python3")
    py = venv_py if os.path.exists(venv_py) else sys.executable
    tasks = ",".join(a.tasks) if a.tasks else "hellaswag,arc_challenge,mmlu"
    # dataset mancanti -> li scarica una volta (fetch_benchmarks.py li mette in --data come JSONL)
    missing=[t for t in tasks.split(",") if not os.path.exists(os.path.join(a.data,f"{t}.jsonl"))]
    if missing:
        print(f"  {C.dim}scarico i dataset mancanti: {', '.join(missing)}{C.r}")
        subprocess.call([py, os.path.join(HERE,"fetch_benchmarks.py"),
                         "--out", a.data, "--tasks", ",".join(missing), "--limit", str(max(a.limit,200))])
    cmd=[py, os.path.join(HERE,"eval_glm.py"), "--snap",a.model,
         "--tasks", tasks, "--limit", str(a.limit), "--data", a.data]
    if a.ram: cmd+=["--ram",str(a.ram)]
    e=dict(os.environ)
    if a.topp: e["TOPP"]=str(a.topp)
    if a.topk: e["TOPK"]=str(a.topk)
    print(f"  {C.dim}decode disk-bound: su hardware lento questo richiede ORE. Alza --limit su macchine veloci.{C.r}\n")
    sys.exit(subprocess.call(cmd, env=e))

def cmd_models(a):
    banner("models")
    reg=models_load()
    if a.action and a.action[0]=="add":
        if len(a.action)<2: sys.exit("uso: coli models add <dir> [nome]")
        d=os.path.abspath(a.action[1])
        if not os.path.isdir(d): sys.exit(f"directory non trovata: {d}")
        name=a.action[2] if len(a.action)>2 else os.path.basename(d)
        reg[name]={"dir":d,"arch":model_arch(d)}
        models_save(reg)
        print(f"  {C.grn}✓{C.r} registrato: {name} · {reg[name]['arch']}\n")
        return
    if not reg: print(f"  {C.dim}nessun modello registrato (coli models add <dir>){C.r}\n"); return
    for n,v in reg.items():
        eng=ENGINES.get(v.get("arch","?"))
        st = "pronto ✓" if eng and os.path.exists(os.path.join(HERE,eng)) else \
             ("motore in sviluppo" if eng else "architettura non supportata")
        ex = "" if os.path.isdir(v.get("dir","")) else " · DIR MANCANTE"
        print(f"   {C.teal}{n:<18}{C.r} {C.gray}{v.get('arch','?'):<14}{C.r} {st}{C.yel}{ex}{C.r}")
        print(f"   {C.dgray}{'':<18} {v.get('dir','')}{C.r}")
    print()

def cmd_convert(a):
    banner("convert")
    # python con torch/safetensors: l'ambiente del progetto se c'e', altrimenti quello corrente
    venv_py=os.path.join(HERE,"mio_env","bin","python3")
    py = venv_py if os.path.exists(venv_py) else sys.executable
    base=[py, os.path.join(HERE,"convert_fp8_to_int4.py"),
          "--repo", a.repo, "--outdir", a.model, "--ebits", str(a.ebits), "--io-bits", str(a.io_bits)]
    if a.xbits: base+=["--xbits",str(a.xbits)]
    # passo 1: modello principale (78 layer). Resumabile: riparte dagli shard mancanti.
    print(f"  {C.dim}[1/2] modello: {' '.join(base)}{C.r}")
    rc=subprocess.call(base)
    if rc!=0: sys.exit(rc)
    if a.no_mtp: sys.exit(0)
    # passo 2: testa MTP (layer 78) -> abilita la decodifica speculativa nativa (piu' veloce)
    print(f"  {C.dim}[2/2] testa MTP (draft speculativi){C.r}")
    sys.exit(subprocess.call(base+["--mtp"]))

def main():
    common=argparse.ArgumentParser(add_help=False)
    common.add_argument("--model", default=DEF_MODEL); common.add_argument("--ram", type=int, default=0)  # 0 = auto (il motore usa l'88% della RAM disponibile)
    common.add_argument("--cap", type=int, default=8); common.add_argument("--ngen", type=int, default=1024)  # rete di sicurezza: la fine vera la decidono gli stop token
    common.add_argument("--topp", type=float, default=0); common.add_argument("--topk", type=int, default=0)
    common.add_argument("--temp", type=float, default=None)  # temperatura token (0=greedy, default 1.0+nucleus .95)
    ap=argparse.ArgumentParser(prog="coli", parents=[common], description="colibrì — GLM-5.2 in locale")
    sub=ap.add_subparsers(dest="cmd")
    sub.add_parser("build", parents=[common]); sub.add_parser("info", parents=[common])
    pr=sub.add_parser("run", parents=[common]);  pr.add_argument("prompt", nargs="*")
    sub.add_parser("chat", parents=[common])
    pm=sub.add_parser("models", parents=[common]); pm.add_argument("action", nargs="*")
    pb=sub.add_parser("bench", parents=[common]); pb.add_argument("tasks", nargs="*")
    pb.add_argument("--limit",type=int,default=40); pb.add_argument("--data",default=os.path.join(HERE,"bench"))
    pc=sub.add_parser("convert", parents=[common]); pc.add_argument("--repo",default="zai-org/GLM-5.2-FP8")
    pc.add_argument("--ebits",type=int,default=4); pc.add_argument("--io-bits",type=int,default=8); pc.add_argument("--xbits",type=int,default=0)
    pc.add_argument("--no-mtp",action="store_true",help="salta la testa MTP (niente draft speculativi)")
    a=ap.parse_args()
    {"build":cmd_build,"info":cmd_info,"run":cmd_run,"chat":cmd_chat,"bench":cmd_bench,
     "convert":cmd_convert,"models":cmd_models}.get(a.cmd, lambda _:(banner(),print(__doc__)))(a)

if __name__=="__main__":
    signal.signal(signal.SIGINT, signal.default_int_handler)
    main()