"""Shared helpers: model loading + pretty printing."""
import os, sys, time
os.environ.setdefault("USE_TF", "0")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

_AGENTS = {}

def get_agent(variant="english"):
    """variant: english | multilingual | typed"""
    if variant in _AGENTS:
        return _AGENTS[variant]
    import laya
    t = time.time()
    if os.path.isdir(variant):
        a = laya.load(variant)
    elif variant == "english":
        # download only the root checkpoint (the other subfolders are several GB each)
        from huggingface_hub import snapshot_download
        path = snapshot_download("convaiinnovations/laya", ignore_patterns=["multilingual/*", "typed-decisions/*"])
        a = laya.load(path)
    elif variant == "multilingual":
        a = laya.load("convaiinnovations/laya", subfolder="multilingual")
    else:
        from huggingface_hub import snapshot_download
        path = snapshot_download("convaiinnovations/laya", allow_patterns=["typed-decisions/*"])
        a = laya.load(path, subfolder="typed-decisions")
    print(f"[laya] loaded {variant} in {time.time()-t:.1f}s", file=sys.stderr)
    if os.environ.get("LAYA_FAST", "1") == "1":
        try:
            sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "kernels"))
            from fast_laya import accelerate
            accelerate(a)
            print("[laya] TileLang fast path enabled (LAYA_FAST=0 to disable)", file=sys.stderr)
        except Exception as e:
            print(f"[laya] fast path unavailable: {e}", file=sys.stderr)
    _AGENTS[variant] = a
    return a

def bar(p, width=20):
    n = int(round(p * width))
    return "█" * n + "░" * (width - n)

def show(result, indent="  "):
    """Pretty-print a laya predict() result."""
    for name, a in result["answers"].items():
        t = a["type"]
        if t == "choice":
            print(f"{indent}{name}: {a['choice']}  (conf {a['confidence']:.2f})")
            for k, v in sorted(a["probabilities"].items(), key=lambda kv: -kv[1]):
                print(f"{indent}    {bar(v)} {v:5.2f}  {k}")
        elif t == "score":
            print(f"{indent}{name}: score={a['score']:.2f}  (conf {a['confidence']:.2f})")
            for k, v in a["probabilities"].items():
                print(f"{indent}    {bar(v)} {v:5.2f}  {a['legend'][k]}")
        else:
            p = a["noul"]
            print(f"{indent}{name}: P(true)={p:.2f}  {bar(p)}")