#!/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 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

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)

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}")
    if not os.path.exists(GLM):
        sys.exit(f"{C.yel}motore non compilato.{C.r} Esegui:  coli build")

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)
    return e

class Spinner:
    FRAMES=["⠋","⠙","⠹","⠸","⠼","⠴","⠦","⠧","⠇","⠏"]
    def __init__(self,label): self.label=label; 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
                sys.stdout.write(f"\r  {C.teal}{self.FRAMES[i%10]}{C.r} {C.dim}{self.label} {el:.0f}s{C.r}\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):
    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([GLM, str(a.cap)], env=e))

def cmd_chat(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([GLM,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)
        extra=" · ".join(l.strip() for l in elog.splitlines() if l.startswith("[RAM_GB") or l.startswith("[PIN]"))
        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}")
        if extra:  print(f"  {C.dgray}{extra}{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
    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
                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}
            sp2=Spinner("pensa…"); 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")
    cmd=[sys.executable, os.path.join(HERE,"eval_glm.py"), "--snap",a.model,
         "--tasks", ",".join(a.tasks) if a.tasks else "hellaswag,arc_challenge,mmlu",
         "--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)
    sys.exit(subprocess.call(cmd, env=e))

def cmd_convert(a):
    banner("convert")
    cmd=[sys.executable, 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: cmd+=["--xbits",str(a.xbits)]
    print(f"  {C.dim}{' '.join(cmd)}{C.r}")
    sys.exit(subprocess.call(cmd))

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)
    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])
    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)
    a=ap.parse_args()
    {"build":cmd_build,"info":cmd_info,"run":cmd_run,"chat":cmd_chat,"bench":cmd_bench,
     "convert":cmd_convert}.get(a.cmd, lambda _:(banner(),print(__doc__)))(a)

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