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

# 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 server_probe(base, api_key=None, timeout=1.5):
    """Is a coli serve alive at `base`? Returns its model_id, or None.
    Probes /health then /v1/models — both cheap, neither touches the engine."""
    import urllib.request, urllib.error
    def get(path):
        req=urllib.request.Request(base.rstrip("/")+path)
        if api_key: req.add_header("Authorization", f"Bearer {api_key}")
        with urllib.request.urlopen(req, timeout=timeout) as r:
            return json.loads(r.read().decode("utf-8","replace"))
    try:
        if get("/health").get("status")!="ok": return None
        data=get("/v1/models").get("data") or []
        return data[0]["id"] if data else None
    except Exception:
        return None

def chat_attached(a, base, model_id):
    """The chat REPL over HTTP against a running `coli serve`.

    Why this exists (the cold-chat cost, measured): spawning a private engine
    pays 34-136 s of resident load on EVERY start, and begins with an empty
    expert cache — hit rate 4% cold vs 55% warm, a ~10x on early decode. A
    resident server pays load once and keeps the LRU warm across sessions;
    its KV slots reuse the conversation prefix, so a continued chat skips
    re-prefill too. The engine byte-protocol stays untouched — this is plain
    OpenAI SSE over localhost, stdlib only."""
    import urllib.request
    print(f"  {C.grn}✦ attached{C.r} {C.dim}to {base} · model {model_id} · the engine stays warm after you quit{C.r}")
    print(f"  {C.dim}type and press Enter · Ctrl-C stops the answer · :reset starts a new conversation · :q exits{C.r}\n")
    msgs=[]
    w=term_w()-4
    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": msgs=[]; print(f"  {C.dim}✦ new conversation{C.r}\n"); continue
        msgs.append({"role":"user","content":msg})
        body=json.dumps({"model":model_id,"messages":msgs,"stream":True,
                         "max_tokens":a.ngen}).encode()
        req=urllib.request.Request(base.rstrip("/")+"/v1/chat/completions", data=body,
                                   headers={"Content-Type":"application/json"})
        if a.api_key: req.add_header("Authorization", f"Bearer {a.api_key}")
        print(f"\n  {C.teal}◆ colibrì{C.r}")
        sp=Spinner("thinking…"); sp.start()
        md=MDStream("  "); reply=[]; first=True; t0=time.time(); interrupted=False
        try:
            with urllib.request.urlopen(req) as r:
                for raw in r:
                    line=raw.decode("utf-8","replace").strip()
                    if not line.startswith("data: "): continue
                    data=line[6:]
                    if data=="[DONE]": break
                    try: ev=json.loads(data)
                    except ValueError: continue
                    for ch in ev.get("choices",[]):
                        d=ch.get("delta",{})
                        txt=d.get("content")
                        if not txt: continue          # ping/ruolo/reasoning: non è testo
                        if first: sp.stop(); first=False
                        md.feed(txt); reply.append(txt)
        except KeyboardInterrupt:
            interrupted=True                          # il server annulla la richiesta alla disconnessione
        except OSError as e:
            sp.stop()
            print(f"\n  {C.yel}[server unreachable: {e}]{C.r}"); break
        sp.stop(); md.close()
        if reply: msgs.append({"role":"assistant","content":"".join(reply)})
        else: msgs.pop()                              # turno vuoto: non sporcare la history
        el=time.time()-t0
        note=" · ⏹ interrupted" if interrupted else ""
        print(f"\r  {C.dgray}└─ ~{len(''.join(reply))//4} tok · {el:.0f}s{note}{C.r}\n")
    print(f"  {C.dim}goodbye — the engine keeps running for the next chat 🐦{C.r}")

def kv_resume_notice(model_dir):
    """SERVE mode silently resumes .coli_kv from disk (glm.c kv_disk_load): a chat
    started today continues a conversation from days ago, with `first=0` so the
    turn is appended WITHOUT the [gMASK]<sop> prefix. The engine does announce it
    on stderr — but nothing here ever shows that: the drain thread's
    p.stderr.read() blocks until EOF, so on a healthy start errlog is still empty
    when the status lines are printed. The warning only appeared once the engine
    DIED, which is exactly when it no longer mattered.

    Measured cost of the silence: a chat inherited 670 tokens of an old Italian
    session ("il mio numero preferito e 7, ricordalo!"). Every later reply came
    back in Italian, and "explain fibonacci in short" was answered about the
    number 7 — the model was being coherent with a context nobody could see, and
    it read as a quantization bug for a day.

    So say it here, in Python, from the file itself: no pipe, no thread, no
    Windows deadlock risk (see the stderr comment below)."""
    p=os.path.join(model_dir, ".coli_kv")
    try:
        with open(p,"rb") as f:
            if f.read(8)!=b"COLIKV1\0": return
            h=struct.unpack("<8i", f.read(32))
        n=h[6]
        if n<1: return
        age=time.time()-os.path.getmtime(p)
        when=f"{age/86400:.0f}d ago" if age>86400 else f"{age/3600:.0f}h ago" if age>3600 else "just now"
        print(f"  {C.yel}↺ resuming a saved conversation: {n} tokens, last written {when}{C.r}")
        print(f"  {C.dgray}  it steers tone, language and topic. :reset clears it · "
              f"KVSAVE=0 disables saving · delete {p} to start clean{C.r}")
    except (OSError, struct.error): pass

def cmd_chat(a):
    # ATTACH: a running `coli serve` beats a private engine every time — the load
    # (34-136 s) and the cache warmth survive between sessions. Explicit --attach
    # wins; otherwise probe localhost quietly and use it if it's there. --no-attach
    # forces the old behaviour. The probe costs ~1 ms when nothing is listening.
    if not getattr(a,"no_attach",False):
        base=getattr(a,"attach",None) or "http://127.0.0.1:8000"
        mid=server_probe(base, getattr(a,"api_key",None))
        if mid:
            banner(f"chat · {mid} · attached")
            chat_attached(a, base, mid); return
        if getattr(a,"attach",None):
            sys.exit(f"--attach: no coli serve answering at {base} (start one with: coli serve --model <dir>)")
    need_model(a.model)
    banner(f"chat · {os.path.basename(a.model)} · ram {a.ram or '-'}GB · topp {a.topp or 'off'}")
    kv_resume_notice(a.model)
    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 serve_pidfile(port): return os.path.join(tempfile.gettempdir(), f"coli-serve-{port}.pid")

def cmd_serve(a):
    need_model(a.model)
    # pidfile: cosi' `coli stop` spegne tutto con un comando, senza pkill a mano.
    # EN: pidfile so `coli stop` can shut everything down without manual pkill.
    try:
        with open(serve_pidfile(a.port),"w") as f: f.write(f"{os.getpid()} {a.model}\n")
    except OSError: pass
    from openai_server import serve
    try:
        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)
    finally:
        try: os.unlink(serve_pidfile(a.port))
        except OSError: pass

def cmd_stop(a):
    """Shut down a running `coli serve` AND its engine — one command, no pkill.
    The engine re-execs itself for OMP tuning, so its process is named `exe`,
    not `glm`: every `pkill -x glm` in history silently killed nothing (that is
    how two 17+5 GB ghost engines OOM'd this box on 2026-07-16). This finds the
    real processes: the pidfile first, then /proc by cmdline/environ — only
    processes that are demonstrably ours (SERVE=1 + our SNAP, or `coli serve`
    in the command line)."""
    banner("stop")
    targets=[]  # (pid, descrizione)
    pf=serve_pidfile(a.port)
    try:
        pid=int(open(pf).read().split()[0])
        os.kill(pid,0); targets.append((pid,f"coli serve (pidfile, port {a.port})"))
    except (OSError,ValueError,IndexError): pass
    for pd in os.listdir("/proc"):
        if not pd.isdigit(): continue
        pid=int(pd)
        try:
            cmd=open(f"/proc/{pd}/cmdline","rb").read().replace(b"\0",b" ").decode("utf-8","replace")
            if "coli" in cmd and " serve" in cmd and pid!=os.getpid():
                if not any(p==pid for p,_ in targets): targets.append((pid,"coli serve (cmdline)"))
            comm=open(f"/proc/{pd}/comm").read().strip()
            if comm in ("glm","exe","olmoe"):
                env=open(f"/proc/{pd}/environ","rb").read().replace(b"\0",b"\n").decode("utf-8","replace")
                if "SERVE=1" in env: targets.append((pid,f"engine `{comm}` (SERVE=1)"))
        except (OSError,PermissionError): continue
    if not targets:
        print(f"  nothing running — no serve on port {a.port}, no SERVE engines"); return
    for pid,desc in targets: print(f"  {'would stop' if a.dry_run else 'stopping'} {pid}: {desc}")
    if a.dry_run: return
    for pid,_ in targets:
        try: os.kill(pid, signal.SIGTERM)
        except OSError: pass
    time.sleep(2.0)
    for pid,_ in targets:
        try: os.kill(pid, signal.SIGKILL); print(f"  {pid}: forced (SIGKILL)")
        except OSError: pass       # gia' morto: bene
    try: os.unlink(pf)
    except OSError: pass
    print(f"  {C.grn}✓ stopped{C.r} — RAM released")

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))])
        # il fetch riprova da solo (#304), ma se l'hub resta giu' il bench gira sui task
        # disponibili invece di passare a eval file inesistenti.
        # EN: the fetch retries on its own (#304), but if the hub stays down the bench
        # runs on the available tasks instead of handing eval nonexistent files.
        still=[t for t in tasks.split(",") if not os.path.exists(os.path.join(a.data,f"{t}.jsonl"))]
        if still:
            tasks=",".join(t for t in tasks.split(",") if t not in still)
            print(f"  {C.yel}skipping (download failed, rerun later): {', '.join(still)}{C.r}")
            if not tasks:
                print(f"  {C.yel}no datasets available — nothing to bench{C.r}"); sys.exit(1)
    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="*")
    pc=sub.add_parser("chat", parents=[common])
    pc.add_argument("--attach", nargs="?", const="http://127.0.0.1:8000", default=None,
                    help="chat against a running `coli serve` instead of spawning an engine "
                         "(keeps the model loaded and the expert cache warm across chat sessions). "
                         "Bare --attach probes localhost:8000.")
    pc.add_argument("--no-attach", action="store_true",
                    help="never auto-attach, always spawn a private engine")
    pc.add_argument("--api-key", default=os.environ.get("COLI_API_KEY"))
    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")))
    pst=sub.add_parser("stop", parents=[common], help="shut down a running coli serve and its engine")
    pst.add_argument("--port",type=int,default=8000); pst.add_argument("--dry-run",action="store_true")
    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,"stop":cmd_stop,"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()
