#!/usr/bin/env python3
"""
colibrì — tiny engine, immense model.
Run GLM-5.2 (744B) locally on CPU with roughly 15-26 GB of RAM.

  coli chat                 interactive chat (loads the model once)
  coli serve                OpenAI-compatible HTTP API (persistent engine)
  coli run "prompt"         one-shot generation
  coli info                 model, RAM, disk, and configuration status
  coli plan                 Disk / RAM / VRAM resource plan
  coli doctor               installation and execution-plan diagnostics
  coli bench [task...]      quality benchmarks (MMLU/HellaSwag/...)
  coli convert              convert GLM-5.2-FP8 to int4, one shard at a time
  coli build                build the engine

Configuration through environment variables or flags (also valid after the subcommand):
  COLI_MODEL=<dir>   model directory (default /home/vincenzo/glm52_i4)
  --ram N            RAM budget in GB (automatically sizes the expert cache)
  --repin N          adapt RAM/VRAM experts every N tokens
  --topp P           adaptive expert top-p             --topk N   fixed top-k
  --ngen N           maximum response tokens           --cap N    cache slots/layer
"""
import os, sys, subprocess, argparse, json, time, signal, shutil, threading, re, codecs, tempfile, textwrap

# The engine mmaps every shard (144+ files); macOS default RLIMIT_NOFILE is 256.
if sys.platform != "win32":
    try:
        import resource
        _soft, _hard = resource.getrlimit(resource.RLIMIT_NOFILE)
        _want = min(65536 if _hard == resource.RLIM_INFINITY else _hard, 65536)
        if _soft < _want:
            resource.setrlimit(resource.RLIMIT_NOFILE, (_want, _hard))
    except (ImportError, ValueError, OSError):
        pass

# Windows: forza output UTF-8 (console cp1252 tronca Unicode box-drawing/emoji)
if sys.platform == "win32":
    for s in (sys.stdout, sys.stderr):
        try: s.reconfigure(encoding="utf-8")
        except (AttributeError, OSError): pass

HERE = os.path.dirname(os.path.abspath(__file__))

# Run-in-place (source checkout, "cd c && ./coli ..."): the engine, the
# support modules (resource_plan.py, doctor.py, openai_server.py) and
# tools/ all live next to this script — unchanged from before.
#
# Installed layout ("make install"): this script is $(PREFIX)/bin/coli,
# while the engine binaries and support files live in
# $(PREFIX)/libexec/colibri, since they aren't meant to be run directly
# by users. COLI_ENGINE overrides the engine path explicitly if neither
# guess is right (e.g. a custom packaging layout).
_EXE = ".exe" if sys.platform == "win32" else ""
_LIBEXEC = os.path.join(os.path.dirname(HERE), "libexec", "colibri")
_here_glm = os.path.join(HERE, "glm" + _EXE)

if os.environ.get("COLI_ENGINE"):
    GLM = os.environ["COLI_ENGINE"]
    TOOLS = os.path.join(os.path.dirname(GLM), "tools")
elif os.path.exists(_here_glm):
    GLM = _here_glm
    TOOLS = os.path.join(HERE, "tools")
else:
    GLM = os.path.join(_LIBEXEC, "glm" + _EXE)
    TOOLS = os.path.join(_LIBEXEC, "tools")
    sys.path.insert(0, _LIBEXEC)   # so `import resource_plan`, `doctor`, `openai_server` still resolve

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}tiny engine, immense model{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}model not found:{C.r} {model}\n  set COLI_MODEL or use --model")
    if not os.path.exists(os.path.join(model,"tokenizer.json")):
        sys.exit(f"{C.yel}tokenizer.json is missing from {model}{C.r}")
    if not os.path.exists(GLM):
        sys.exit(f"{C.yel}engine is not built.{C.r} Run: coli build")

def cuda_binary():
    if not os.path.exists(GLM) or sys.platform != "linux": return False
    try:
        linked=subprocess.run(["ldd",GLM],capture_output=True,text=True,timeout=3)
        return any("libcudart" in line and "not found" not in line
                   for line in linked.stdout.splitlines())
    except (OSError,subprocess.SubprocessError): return False

def resource_request(a, env):
    ctx=a.ctx or int(env.get("CTX",4096))
    ram=a.ram or float(env.get("RAM_GB",0))
    vram=a.vram or float(env.get("CUDA_EXPERT_GB",0))
    gpu=a.gpu
    if gpu is None:
        gpu=env.get("COLI_GPUS",env.get("COLI_GPU","auto"))
    devices=None if gpu=="auto" else ([] if gpu=="none" else
        [int(value) for value in gpu.split(",")])
    return ram,ctx,devices,vram

def env_for(a):
    e = dict(os.environ, SNAP=a.model)
    if sys.platform == "win32":
        # COLI_NO_OMP_TUNE spegne SOLO il blocco OMP (stesso perimetro del
        # self-exec di glm.c; presence-based come nel motore: impostarla a
        # qualsiasi valore, anche 0, disattiva). I default I/O piu' sotto
        # restano attivi: kill-switch dedicati = le var stesse (DIRECT=0 ecc.)
        if not e.get("COLI_NO_OMP_TUNE"):
            # parita' col tuning OMP self-exec di glm.c (solo Linux/FreeBSD, e
            # comunque saltato sotto CUDA/Metal): libgomp legge queste variabili
            # prima di main, quindi su Windows vanno nell'ambiente del figlio.
            # niente OMP_PROC_BIND/OMP_PLACES: la libgomp di MinGW non supporta
            # l'affinity su Windows ("Affinity not supported on this configuration")
            from resource_plan import physical_cpu_count
            for k, v in (("OMP_WAIT_POLICY", "active"),
                         ("GOMP_SPINCOUNT", "200000"),
                         ("OMP_DYNAMIC", "FALSE"),
                         ("OMP_NUM_THREADS", str(physical_cpu_count()))):
                e.setdefault(k, v)
        # Default Windows misurati sul box di riferimento (docs/tuning-9950x3d-5090.md),
        # tutti lossless e tutti setdefault (un override esplicito vince sempre):
        # - DIRECT=1: 10.7 GB/s O_DIRECT vs 9.0 buffered (iobench, anche a cache
        #   calda); nel motore 0.48 -> 1.02 tok/s. Upstream #162: 1.47x.
        # - PIPE=1: overlap load/matmul, +8% sopra DIRECT (byte-identico, riordina
        #   solo l'I/O). PIPE_WORKERS resta al default 8 (sweep 4/8/16 piatto).
        # - PILOT_REAL=1: prefetch cross-layer con load veri (unico prefetch
        #   funzionante su Windows: fadvise e' no-op), +11%, hit rate +19 punti.
        e.setdefault("DIRECT", "1")
        e.setdefault("PIPE", "1")
        e.setdefault("PILOT_REAL", "1")
    e["COLI_POLICY"]=a.policy
    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
    if a.repin: e["REPIN"]=str(a.repin)
    if a.ctx: e["CTX"]=str(a.ctx)
    if a.auto_tier:
        from resource_plan import build_plan, environment_for_plan, format_bytes
        if a.gpu is not None:
            e.pop("COLI_GPU",None); e.pop("COLI_GPUS",None)
            if a.gpu=="none":
                e["COLI_CUDA"]="0"; e.pop("CUDA_EXPERT_GB",None); e.pop("CUDA_DENSE",None)
            else: e.pop("COLI_CUDA",None)
        elif e.get("COLI_CUDA")=="0":
            e.pop("COLI_GPU",None); e.pop("COLI_GPUS",None)
            e.pop("CUDA_EXPERT_GB",None); e.pop("CUDA_DENSE",None)
        if a.vram and a.gpu!="none": e["CUDA_EXPERT_GB"]=str(a.vram)
        try:
            ram,ctx,devices,vram=resource_request(a,e)
            plan=build_plan(a.model,ram,ctx,devices,vram,policy=a.policy)
        except (OSError,ValueError,json.JSONDecodeError) as error:
            sys.exit(f"{C.yel}invalid resource plan:{C.r} {error}")
        has_cuda=cuda_binary()
        e=environment_for_plan(plan,e,has_cuda)
        rt=plan["tiers"]["ram"]; vt=plan["tiers"]["vram"]
        gpu=f" · VRAM {format_bytes(vt['budget_bytes'])}" if has_cuda and vt["devices"] else " · CPU"
        print(f"  {C.dim}[PLAN] RAM {format_bytes(rt['budget_bytes'])} · cap {rt['cache_slots_per_layer']}/layer{gpu}{C.r}",file=sys.stderr)
    else:
        # --gpu/--vram SENZA --auto-tier: prima venivano ignorati in silenzio e il run
        # partiva CPU-only senza alcun avviso — benchmark "GPU" pubblicati per errore (#121).
        if a.gpu is not None:
            e.pop("COLI_GPU",None); e.pop("COLI_GPUS",None)
            if a.gpu=="none":
                e["COLI_CUDA"]="0"; e.pop("CUDA_EXPERT_GB",None); e.pop("CUDA_DENSE",None)
            else:
                if not cuda_binary():
                    sys.exit(f"{C.yel}--gpu needs the CUDA build:{C.r} make glm CUDA=1 (this binary is CPU-only)")
                e["COLI_CUDA"]="1"
                if a.gpu!="auto": e["COLI_GPUS"]=a.gpu
                e.setdefault("CUDA_DENSE","1")
        if a.vram and a.gpu!="none":
            if not cuda_binary():
                sys.exit(f"{C.yel}--vram needs the CUDA build:{C.r} make glm CUDA=1 (this binary is CPU-only)")
            e["COLI_CUDA"]="1"; e["CUDA_EXPERT_GB"]=str(a.vram)
    return e

# ---------- rendering markdown in STREAMING per il terminale ----------
class MDStream:
    """Interpreta il markdown della risposta mentre arriva: i ``` diventano riquadri,
    **x** grassetto vero, `x` colorato, # titoli, - puntini. I marker non si vedono mai.
    Regge i chunk spezzati a meta' marker (hold-back) e l'output sporco (``` doppi)."""
    def __init__(self, indent="  "):
        self.ind=indent
        self.cur=""                      # riga parziale non ancora emessa
        self.code=False; self.lang=""
        self.bold=False; self.icode=False
        self.justclosed=False            # l'ultima riga era una chiusura ```? (anti ``` doppi)
        self.printed=0                   # caratteri della riga corrente gia' emessi
    def _fence(self, line):
        lang=line.strip()[3:].strip().strip("`")
        if not self.code:
            if not lang and self.justclosed: return   # ``` orfano dopo una chiusura: rumore, ignora
            self.code=True; self.lang=lang
            sys.stdout.write(f"{self.ind}{C.dgray}\u256d\u2500 {lang or 'code'}{C.r}\n")
        elif lang:                       # ```lang mentre siamo GIA' in code: chiudi e riapri
            sys.stdout.write(f"{self.ind}{C.dgray}\u2570\u2500{C.r}\n{self.ind}{C.dgray}\u256d\u2500 {lang}{C.r}\n")
            self.lang=lang
        else:
            self.code=False; self.justclosed=True
            sys.stdout.write(f"{self.ind}{C.dgray}\u2570\u2500{C.r}\n")
    def _inline(self, txt, out):
        i=0
        while i<len(txt):
            ch=txt[i]
            if ch=="`":
                self.icode=not self.icode
                out.append(C.org if self.icode else C.r); i+=1; continue
            if ch=="*":
                j=i
                while j<len(txt) and txt[j]=="*": j+=1
                if j-i>=2:               # **/***: grassetto on/off, gli asterischi spariscono
                    self.bold=not self.bold
                    out.append(C.b if self.bold else C.r)
                else: out.append("*")    # * singolo: lascialo (moltiplicazioni ecc.)
                i=j; continue
            out.append(ch); i+=1
    def _line(self, line, partial=False):
        if not partial and line.lstrip().startswith("```"):
            self._fence(line); self.printed=0; return
        if line.strip(): self.justclosed=False
        seg=line[self.printed:]          # emetti solo la parte nuova della riga
        out=[]
        if self.code:
            if self.printed==0: out.append(f"{self.ind}{C.dgray}\u2502{C.r} {C.cyan}")
            out.append(seg)
        else:
            if self.printed==0:
                out.append(self.ind)
                st=seg.lstrip()
                if st.startswith("#"):           # titolo: via i #, grassetto teal
                    seg=st.lstrip("#").strip(); out.append(f"{C.teal}{C.b}"); self.bold=True
                elif st.startswith(("- ","* ")): # lista: puntino vero
                    seg=st[2:]; out.append(f"{C.teal}\u2022{C.r} ")
            self._inline(seg,out)
        sys.stdout.write("".join(out)); sys.stdout.flush()
        self.printed=len(line)
        if not partial:                  # fine riga: reset stati inline (robusto ai marker orfani)
            sys.stdout.write(C.r+"\n"); sys.stdout.flush()
            self.bold=self.icode=False
            self.printed=0
    def feed(self, s):
        self.cur+=s
        while "\n" in self.cur:
            line,self.cur=self.cur.split("\n",1)
            self._line(line)
        st=self.cur.lstrip()             # riga parziale: possibile fence? aspetta il newline
        if st and (st.startswith("```") or (len(st)<3 and "```".startswith(st))):
            return
        if st.startswith("#") and self.printed==0:
            return                        # titolo: rendi la riga intera al newline
        hold=0                            # trattieni marker potenzialmente spezzati in coda
        while hold<len(self.cur) and self.cur[-1-hold] in "*`": hold+=1
        safe=self.cur[:len(self.cur)-hold] if hold else self.cur
        if len(safe)>self.printed: self._line(safe, partial=True)
    def close(self):
        if self.cur: self._line(self.cur); self.cur=""
        if self.code:
            sys.stdout.write(f"\n{self.ind}{C.dgray}\u2570\u2500{C.r}"); self.code=False
        sys.stdout.write(C.r); sys.stdout.flush()

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 engine_diag(p, errlog=None):
    """Lo stdout del motore ha chiuso: il processo e' morto. Dire PERCHE'.
    Un SIGKILL dell'OOM-killer del kernel e' altrimenti invisibile — niente errore,
    niente exit code, il motore muore muto e sembra un bug nostro (issue #305).
    Quasi sempre e' memoria: il picco reale ha sforato la RAM della macchina."""
    try: rc=p.wait(timeout=5)
    except Exception: rc=p.poll()
    if rc is None:                      why="its output closed but the process is still alive"
    elif rc<0:                                        # POSIX: morte per segnale
        try: why=f"killed by {signal.Signals(-rc).name}"
        except Exception: why=f"killed by signal {-rc}"
    elif rc>0:                          why=f"exit code {rc}"
    else:                               why="exited cleanly"
    print(f"\n  {C.yel}[engine terminated: {why}]{C.r}")
    if rc is not None and rc<0 and -rc==getattr(signal,"SIGKILL",-1):
        print(f"  {C.yel}nothing in the engine sends SIGKILL to itself: this is the kernel's\n"
              f"  OOM-killer. The peak RSS exceeded the machine's free memory.\n"
              f"  Lower --ram, lower PIN_GB, or shorten the context.{C.r}")
    if errlog is not None:
        try:
            errlog.flush()
            for _ in range(20):         # il drain thread scrive su errlog quando il child chiude
                tail=open(errlog.name).read().strip()
                if tail: break
                time.sleep(0.05)
            if tail: print(f"  {C.dgray}" + "\n  ".join(tail.splitlines()[-6:]) + f"{C.r}")
        except Exception: pass

def stream_turn(p, sentinel, on_bytes):
    """legge fino alla sentinella; on_bytes riceve i chunk della risposta. Poi legge la riga STAT.
    Il PRIMO Ctrl-C durante lo stream non chiude la sessione: il motore (handler SIGINT)
    chiude il turno per la via del tetto NGEN e noi dreniamo fino alla sentinella.
    Un SECONDO Ctrl-C esce davvero."""
    pend=b""; interrupted=False
    while True:
        try:
            b=p.stdout.read(1)
        except KeyboardInterrupt:
            if interrupted or p.poll() is not None: raise
            interrupted=True
            try: p.send_signal(signal.SIGINT)   # non-TTY: il motore potrebbe non aver visto il Ctrl-C
            except Exception: pass
            print(f"\n  {C.yel}⏹ stopping… (Ctrl-C again to quit){C.r}", flush=True)
            continue
        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)
            st={"tok":int(m.group(1)),"tps":float(m.group(2)),"hit":float(m.group(3)),"rss":float(m.group(4))} if m else {}
            if interrupted: st["interrupted"]=True
            return st
        if len(pend)>len(sentinel):
            out=pend[:-len(sentinel)]; pend=pend[-len(sentinel):]
            on_bytes(out)

# ---------- comandi ----------
def cmd_build(a):
    banner("build")
    if not os.path.exists(os.path.join(HERE, "Makefile")):
        sys.exit(f"{C.yel}coli build{C.r} only works from a source checkout (this is an installed copy).\n"
                  f"  Clone https://github.com/JustVugg/colibri and run ./setup.sh, or make -C c glm.")
    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("model", 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("shards", f"{len(sts)} files · {sz/1e9:.0f} GB on disk")
    else:
        print(f"   {C.yel}config.json is missing (incomplete conversion?){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 total · {av:.1f} GB available")
    except Exception: pass
    try:
        fs = shutil.disk_usage(a.model if os.path.isdir(a.model) else HERE)
        row("disk", f"{fs.free/1e9:.0f} GB free")
    except OSError:
        row("disk", "? GB (unavailable)")
    row("engine", "ready ✓" if os.path.exists(GLM) else "not built (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_plan(a):
    from resource_plan import build_plan, format_plan
    try:
        ram,ctx,devices,vram=resource_request(a,os.environ)
        if ctx<1: raise ValueError("--ctx must be positive")
        if a.vram<0: raise ValueError("--vram cannot be negative")
        plan=build_plan(a.model,ram,ctx,devices,vram,policy=a.policy)
    except (OSError, ValueError, json.JSONDecodeError) as error:
        sys.exit(f"{C.yel}cannot create resource plan:{C.r} {error}")
    if a.json:
        print(json.dumps(plan,indent=2))
        return
    banner("plan · Disk / RAM / VRAM")
    print(textwrap.indent(format_plan(plan),"  "))
    print()

def cmd_doctor(a):
    from doctor import exit_code, format_doctor, run_doctor
    try:
        ram,ctx,devices,vram=resource_request(a,os.environ)
        if ctx<1: raise ValueError("--ctx must be positive")
        if ram<0: raise ValueError("--ram cannot be negative")
        if vram<0: raise ValueError("--vram cannot be negative")
    except ValueError as error:
        report={"schema_version":1,"status":"error","model":os.path.abspath(a.model),
                "checks":[{"id":"config.arguments","status":"fail","summary":str(error)}],
                "plan":None}
        print(json.dumps(report,indent=2) if a.json else format_doctor(report))
        return 2
    report=run_doctor(a.model,ram,ctx,devices,vram,engine_path=GLM)
    print(json.dumps(report,indent=2) if a.json else format_doctor(report))
    return exit_code(report)

def cmd_run(a):
    need_model(a.model)
    prompt=" ".join(a.prompt) if a.prompt else sys.exit('usage: coli run "your 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"
    # stderr -> PIPE, NOT stderr=errlog (file). On Windows/MinGW, pointing the
    # child's stderr at a file/DEVNULL handle stalls the CRT so stdout (the byte
    # protocol coli reads one byte at a time) never flushes and chat hangs at
    # ~10 GB resident. A PIPE whose read end nobody drains still works: the
    # engine emits only ~400 bytes of status to stderr, which fits comfortably
    # in the OS pipe buffer, so it never blocks. We snapshot stderr into errlog
    # once the READY sentinel arrives, so the status-line display below works
    # exactly as before. (Do NOT add a concurrent stderr drain thread: on
    # Windows, reading two child pipes simultaneously deadlocks CPython's IO.)
    p=subprocess.Popen([GLM,str(a.cap)], env=e, stdin=subprocess.PIPE,
                       stdout=subprocess.PIPE, stderr=subprocess.PIPE, bufsize=0)
    sp=Spinner("waking the giant (744B)…"); sp.start()
    st=stream_turn(p, READY, lambda b: None)
    sp.stop()
    if st is None:
        try: errlog.write(p.stderr.read().decode("utf-8","replace"))
        except (OSError, ValueError): pass
        errlog.flush(); errlog.seek(0); print(errlog.read()[-1500:])
        engine_diag(p)        # perche' e' morto (OOM-kill compreso); errlog e' gia' stampato sopra
        sys.exit("the engine exited while loading")
    p.stdout.readline()   # TIERS line (web-dashboard protocol): emitted once right after STAT;
                          # left unread it leaks into the first answer's text
    # READY received. Drain the child's stderr into errlog without blocking:
    # the engine is still alive (blocked on stdin), so a plain read() would
    # hang forever waiting for EOF. A short bounded drain grabs the ~400 bytes
    # of load-time status ([RAM_GB], [MTP], ...) that were already emitted.
    _drain_box={"done":False}
    def _drain():
        try: errlog.write(p.stderr.read().decode("utf-8","replace"))
        except (OSError, ValueError): pass
        _drain_box["done"]=True
    threading.Thread(target=_drain, daemon=True).start()
    _drain_box["th"]=threading.current_thread()
    for _ in range(20):           # up to ~1s for the load-status lines
        if _drain_box["done"]: break
        time.sleep(0.05)
    errlog.flush()
    try:
        elog=open(errlog.name).read()
        mload=re.search(r"loaded in ([0-9.]+)s \| resident dense: ([0-9.]+) MB", elog)
        if mload: print(f"  {C.grn}✓{C.r} ready in {mload.group(1)}s {C.dim}· resident {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]","[KV]")):
                l=re.sub(r" ?\(?/[^ )]+\)?","",l.strip())       # via i percorsi lunghi
                l=re.sub(r" from$","",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}type and press Enter · Ctrl-C stops the answer · :more continues · :reset clears memory · :q exits{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}✦ memory cleared{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("thinking…", tick=prefill_tick); sp2.start()
            md=MDStream("  ")            # markdown -> terminale, in streaming
            raw=os.environ.get("COLI_RAW")=="1"
            def echo(bs, _dec=dec, _st=state):
                if _st["first"]:
                    sp2.stop(); _st["first"]=False
                    if raw: sys.stdout.write("  ")
                s=_dec.decode(bs)
                if not s: return
                if raw: sys.stdout.write(s.replace("\n","\n  ")); sys.stdout.flush()
                else: md.feed(s)
            t0=time.time()
            st=stream_turn(p, END, echo)
            if not raw: md.close()
            sp2.stop()
            if st is None: engine_diag(p, errlog); 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.get("interrupted"):
                    print(f"  {C.yel}⏹ interrupted; type :more to continue the response{C.r}")
                elif st["tok"]>=a.ngen:
                    print(f"  {C.yel}…stopped at --ngen ({a.ngen}); type :more to continue the response{C.r}")
                print()
            else:
                if st.get("interrupted"): print(f"  {C.yel}⏹ interrupted{C.r}")
                print()
    except KeyboardInterrupt:
        print(f"\n  {C.dim}interrupted{C.r}")
    finally:
        try: p.stdin.close(); p.terminate()
        except Exception: pass
        try: os.unlink(errlog.name)
        except Exception: pass
    print(f"  {C.teal}goodbye{C.r} {C.dim}— the hummingbird returns to its nest{C.r} 🐦\n")

def cmd_serve(a):
    need_model(a.model)
    from openai_server import serve
    serve(a.model, a.host, a.port, a.model_id, a.api_key,
          a.cap,a.ngen,GLM,env_for(a),a.cors_origin,
          a.max_queue,a.queue_timeout,a.kv_slots)

def cmd_web(a):
    """serve + open the dashboard in the browser once the API answers."""
    need_model(a.model)
    dist = os.path.join(os.path.dirname(os.path.abspath(HERE)), "web", "dist")
    if not os.path.exists(os.path.join(dist, "index.html")):
        print(f"{C.yel}web UI not built:{C.r} run  cd web && npm install && npm run build  first;")
        print("serving the API anyway (the dashboard will 404 until built).")
    url = f"http://{a.host}:{a.port}/"
    if not getattr(a, "no_browser", False):
        import threading, urllib.request, webbrowser
        def opener():
            for _ in range(600):                    # the 744B engine takes minutes to load
                time.sleep(2)
                try:
                    urllib.request.urlopen(f"http://{a.host}:{a.port}/health", timeout=2)
                    webbrowser.open(url); return
                except OSError:
                    continue
        threading.Thread(target=opener, daemon=True).start()
    print(f"dashboard: {url}  (opens automatically when the engine is ready)")
    cmd_serve(a)

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", "Scripts" if sys.platform == "win32" else "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}downloading missing datasets: {', '.join(missing)}{C.r}")
        subprocess.call([py, os.path.join(TOOLS,"fetch_benchmarks.py"),
                         "--out", a.data, "--tasks", ",".join(missing), "--limit", str(max(a.limit,200))])
    cmd=[py, os.path.join(TOOLS,"eval_glm.py"), "--glm", GLM, "--snap",a.model,
         "--tasks", tasks, "--limit", str(a.limit), "--data", a.data]
    if a.ram: cmd+=["--ram",str(a.ram)]
    e=env_for(a)
    print(f"  {C.dim}decode is disk-bound: this takes HOURS on slow hardware. Raise --limit on faster machines.{C.r}\n")
    sys.exit(subprocess.call(cmd, env=e))

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", "Scripts" if sys.platform == "win32" else "bin", "python3")
    py = venv_py if os.path.exists(venv_py) else sys.executable
    base=[py, os.path.join(TOOLS,"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] model: {' '.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). SEMPRE int8: a int4 i draft sbagliano quasi sempre
    # (acceptance 0-4% vs 39-59%, misurato — issue #8) e la speculazione non parte mai.
    mtp_cmd=list(base); i=mtp_cmd.index("--ebits"); mtp_cmd[i+1]=str(max(8,a.ebits))
    print(f"  {C.dim}[2/2] int8 MTP head (speculative drafts){C.r}")
    sys.exit(subprocess.call(mtp_cmd+["--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("--auto-tier",action="store_true",help="automatically apply the RAM/VRAM plan")
    common.add_argument("--ctx",type=int,default=0)
    common.add_argument("--gpu",default=None,help="auto, none, or a device list such as 0,1")
    common.add_argument("--vram",type=float,default=0,help="total VRAM budget in GB (0=auto)")
    common.add_argument("--policy",choices=("quality","balanced","experimental-fast"),
                        default=os.environ.get("COLI_POLICY","quality"),
                        help="resource policy (explicit --topk/--topp overrides warn and proceed)")
    common.add_argument("--repin", type=int, default=0, help="adapt RAM/VRAM experts every N tokens")
    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ì — run GLM-5.2 locally")
    sub=ap.add_subparsers(dest="cmd")
    sub.add_parser("build", parents=[common]); sub.add_parser("info", parents=[common])
    pp=sub.add_parser("plan",parents=[common])
    pp.add_argument("--json",action="store_true")
    pd=sub.add_parser("doctor",parents=[common])
    pd.add_argument("--json",action="store_true",help="emit a versioned JSON report")
    pr=sub.add_parser("run", parents=[common]);  pr.add_argument("prompt", nargs="*")
    sub.add_parser("chat", parents=[common])
    ps=sub.add_parser("serve", parents=[common])
    ps.add_argument("--host",default="127.0.0.1"); ps.add_argument("--port",type=int,default=8000)
    ps.add_argument("--model-id",default=os.environ.get("COLI_MODEL_ID","glm-5.2-colibri"))
    ps.add_argument("--api-key",default=os.environ.get("COLI_API_KEY"))
    ps.add_argument("--cors-origin",action="append",default=None)
    ps.add_argument("--max-queue",type=int,default=int(os.environ.get("COLI_MAX_QUEUE","8")))
    ps.add_argument("--queue-timeout",type=float,default=float(os.environ.get("COLI_QUEUE_TIMEOUT","300")))
    ps.add_argument("--kv-slots",type=int,default=int(os.environ.get("COLI_KV_SLOTS","1")))
    pw=sub.add_parser("web", parents=[common], help="serve + open the dashboard in a browser")
    for arg,kw in (("--host",dict(default="127.0.0.1")),("--port",dict(type=int,default=8000)),
                   ("--model-id",dict(default=os.environ.get("COLI_MODEL_ID","glm-5.2-colibri"))),
                   ("--api-key",dict(default=os.environ.get("COLI_API_KEY"))),
                   ("--cors-origin",dict(action="append",default=None)),
                   ("--max-queue",dict(type=int,default=int(os.environ.get("COLI_MAX_QUEUE","8")))),
                   ("--queue-timeout",dict(type=float,default=float(os.environ.get("COLI_QUEUE_TIMEOUT","300")))),
                   ("--kv-slots",dict(type=int,default=int(os.environ.get("COLI_KV_SLOTS","1"))))):
        pw.add_argument(arg,**kw)
    pw.add_argument("--no-browser",action="store_true",help="don't auto-open the browser")
    pb=sub.add_parser("bench", parents=[common]); pb.add_argument("tasks", nargs="*")
    pb.add_argument("--limit",type=int,default=40)
    if sys.platform == "win32":
        _cache_root = os.environ.get("LOCALAPPDATA", os.path.expanduser("~\\AppData\\Local"))
    else:
        _cache_root = os.environ.get("XDG_CACHE_HOME", os.path.expanduser("~/.cache"))
    _bench_cache = os.path.join(_cache_root, "colibri", "bench")
    pb.add_argument("--data",default=_bench_cache)
    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="skip the MTP head (no speculative drafts)")
    a=ap.parse_args()
    handler={"build":cmd_build,"info":cmd_info,"plan":cmd_plan,"doctor":cmd_doctor,
             "run":cmd_run,"chat":cmd_chat,"serve":cmd_serve,"bench":cmd_bench,
             "convert":cmd_convert,"web":cmd_web}.get(a.cmd)
    if handler: sys.exit(handler(a) or 0)
    banner(); print(__doc__)

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