"""
llm_local.py — local sub-agent pool (llama.cpp runtime, no cloud APIs).

Two independent model instances ("sub-agents"), each with its own lock, so
features never block each other with "model busy":

    fast  · Summary Sub-Agent (Chan-Tuned Qwen3-1.7B)
            → Explain-in-English, sector-rotation narrative, news briefs.
            Small = quick CPU prefill, answers start streaming in seconds.
    deep  · Analyst — Qwen3-4B Q4_K_M by default (swappable in the Model tab)
            → the multi-step Auto Research agent's report writing.

Both run through llama.cpp (llama-cpp-python) and are far below the 32B cap;
the fast worker doubles as the "Tiny Titan" (≤4B) story. ~5 GB RAM total on a
32 GB Space. Earns "Off the Grid" + "Llama Champion".
"""
from __future__ import annotations

import os
import re
import threading

import paths  # sets HF_HOME + sys.path for /data persistence
from huggingface_hub import hf_hub_download

# name -> (HF repo, gguf filename)
MODEL_ZOO = {
    "Chan-Tuned Qwen3-1.7B · my fine-tune": (
        "ranranrunforit/chan-compass-qwen3-1.7b-gguf", "qwen3-1.7b.Q8_0.gguf"),
    "Qwen3-1.7B · Tiny Titan (≤4B award class)": (
        "Qwen/Qwen3-1.7B-GGUF", "Qwen3-1.7B-Q8_0.gguf"),
    "Qwen3-4B  · default — fast + smart, still ≤4B": (
        "Qwen/Qwen3-4B-GGUF", "Qwen3-4B-Q4_K_M.gguf"),
    "Qwen3-8B  · best balance on 8 vCPU / 32 GB": (
        "Qwen/Qwen3-8B-GGUF", "Qwen3-8B-Q4_K_M.gguf"),
    "Qwen3-14B · max quality (still far under 32B cap)": (
        "Qwen/Qwen3-14B-GGUF", "Qwen3-14B-Q4_K_M.gguf"),
}
# Only the Summary sub-agent (Signals · Explain) uses the published
# fine-tune; every other sub-agent stays on the stock models.
FAST_MODEL = "Qwen3-1.7B · Tiny Titan (≤4B award class)"
TRANSLATOR_MODEL = "Chan-Tuned Qwen3-1.7B · my fine-tune"
DEFAULT_MODEL = "Qwen3-4B  · default — fast + smart, still ≤4B"

_THINK_RE = re.compile(r"<think>.*?</think>", re.S)
_NCPU = max(2, (os.cpu_count() or 4))

# One dedicated sub-agent per feature — independent locks, so Signals-Explain,
# Rotation narrative, News briefs and Auto-Research never fight over a model.
# Three tiny 1.7B instances share ONE GGUF file on disk (~2 GB RAM each) and
# the 4B Analyst writes reports. Total ≈ 9 GB on a 32 GB Space.
WORKER_LABEL = {
    "interpreter": "Interpreter sub-agent (Signals · Explain)",
    "narrator":   "Narrator sub-agent (Sector Rotation)",
    "reporter":   "Reporter sub-agent (News · Research support)",
    "analyst":    "Analyst sub-agent (Auto Research)",
}


def _mk(model):
    return {"model": model, "llm": None, "lock": threading.Lock(),
            "load_lock": threading.Lock(), "stage": "idle", "detail": "", "ts": None}


WORKERS = {
    "interpreter": _mk(TRANSLATOR_MODEL),
    "narrator":   _mk(FAST_MODEL),
    "reporter":   _mk(FAST_MODEL),
    "analyst":    _mk(DEFAULT_MODEL),
}
# legacy aliases
_ALIAS = {"fast": "interpreter", "deep": "analyst"}


def _wk(worker: str) -> str:
    return _ALIAS.get(worker, worker)

_install_lock = threading.Lock()


def _set_stage(worker: str, stage: str, detail: str = ""):
    import datetime as _dt
    w = WORKERS[worker]
    w.update(stage=stage, detail=detail[:400],
             ts=_dt.datetime.utcnow().strftime("%H:%M:%S UTC"))
    try:
        import automation
        automation._log(f"[{WORKER_LABEL[worker]}] {stage}: {detail[:140]}")
    except Exception:
        pass


# ─────────────────────── runtime install (once, persisted) ───────────────────────
# Installed at RUNTIME, not at Space build time: the HF build container has
# little RAM and gets OOM-killed compiling the C++ extension; the runtime
# container has the real hardware. Prebuilt CPU wheel first, capped-parallelism
# source build as fallback. Persisted to /data/pylibs.
_WHEEL_INDEX = "https://abetlen.github.io/llama-cpp-python/whl/cpu"
_LLAMA_REQ = "llama-cpp-python>=0.3.8"   # >=0.3.8 → Qwen3 architecture support


def _ensure_llama_cpp(worker: str) -> str:
    try:
        import llama_cpp  # noqa: F401
        return ""
    except ImportError:
        pass
    import subprocess
    import sys
    with _install_lock:
        try:  # another thread may have finished it while we waited
            import llama_cpp  # noqa: F401
            return ""
        except ImportError:
            pass
        env = dict(os.environ)
        env["CMAKE_BUILD_PARALLEL_LEVEL"] = "4"
        _set_stage(worker, "installing llama.cpp runtime",
                   "trying official prebuilt CPU wheel (≈1 min)…")
        if paths.PERSISTENT:
            base = [sys.executable, "-m", "pip", "install", "--prefer-binary",
                    "--target", paths.PYLIBS_DIR]
        else:
            base = [sys.executable, "-m", "pip", "install", "--user", "--prefer-binary"]
        r = subprocess.run(base + ["--extra-index-url", _WHEEL_INDEX,
                                   "--only-binary", "llama-cpp-python", _LLAMA_REQ],
                           capture_output=True, text=True, env=env, timeout=600)
        if r.returncode != 0:
            _set_stage(worker, "installing llama.cpp runtime",
                       "no prebuilt wheel matched — compiling from source "
                       "(one-time ~10-15 min; other tabs keep working)…")
            r = subprocess.run(base + ["--extra-index-url", _WHEEL_INDEX, _LLAMA_REQ],
                               capture_output=True, text=True, env=env, timeout=2400)
        if r.returncode != 0:
            err = (r.stderr or r.stdout or "")[-800:]
            _set_stage(worker, "install FAILED", err)
            return "Could not install llama-cpp-python at runtime:\n" + err
        import importlib
        import site
        cands = [paths.PYLIBS_DIR] if paths.PERSISTENT else []
        usp = site.getusersitepackages()
        cands += usp if isinstance(usp, list) else [usp]
        for p in cands:
            if p and p not in sys.path:
                sys.path.append(p)
        importlib.invalidate_caches()
        try:
            import llama_cpp  # noqa: F401
            return ""
        except Exception as e:
            _set_stage(worker, "install FAILED", f"installed but import failed: {e}")
            return f"Installed but import failed: {e}"


# ─────────────────────── loading ───────────────────────
def load_model(name: str, worker: str = "analyst") -> str:
    worker = _wk(worker)
    """Load a GGUF into a worker slot. Non-blocking: if that worker is already
    installing/loading, returns its live stage instead of hanging the click."""
    w = WORKERS[worker]
    if not w["load_lock"].acquire(timeout=2):
        return (f"⏳ {WORKER_LABEL[worker]} is busy — current stage: "
                f"**{w['stage']}** ({w['detail'] or '…'}). Press “↻ Refresh status”.")
    try:
        if w["llm"] is not None and w["model"] == name:
            return f"Already loaded on {WORKER_LABEL[worker]}: {name}"
        err = _ensure_llama_cpp(worker)
        if err:
            return err
        try:
            from llama_cpp import Llama
        except Exception as e:
            _set_stage(worker, "import FAILED", str(e))
            return f"llama-cpp-python is not available: {e}"
        repo, fname = MODEL_ZOO[name]
        try:
            _set_stage(worker, "downloading GGUF",
                       f"{repo}/{fname} (cached on /data after first time)")
            path = hf_hub_download(repo_id=repo, filename=fname)
        except Exception as e:
            _set_stage(worker, "download FAILED", str(e))
            return f"Could not download {repo}/{fname}: {e}"
        try:
            _set_stage(worker, "loading model into RAM", name)
            w["llm"] = None
            small = worker != "analyst"
            w["llm"] = Llama(
                model_path=path,
                n_ctx=4096 if small else 6144,
                n_threads=(4 if small else _NCPU),   # leave headroom for parallel agents
                n_threads_batch=(6 if small else _NCPU),
                n_batch=512,
                verbose=False,
            )
            w["model"] = name
            _set_stage(worker, "ready", name)
            return f"✅ {WORKER_LABEL[worker]} ready: {name}"
        except Exception as e:
            w["llm"] = None
            _set_stage(worker, "load FAILED", str(e))
            return f"Failed to load model: {e}"
    finally:
        w["load_lock"].release()


def auto_load_all():
    """Startup: tiny agents first (one small GGUF download serves all three),
    then the Analyst. Runs in a background thread."""
    for key in ("interpreter", "narrator", "reporter", "analyst"):
        load_model(WORKERS[key]["model"], worker=key)


# ─────────────────────── status ───────────────────────
def is_loaded(worker: str = None) -> bool:
    if worker:
        return WORKERS[_wk(worker)]["llm"] is not None
    return any(w["llm"] is not None for w in WORKERS.values())


def status() -> str:
    lines = []
    for key in ("interpreter", "narrator", "reporter", "analyst"):
        w = WORKERS[key]
        label = WORKER_LABEL[key]
        if w["llm"] is not None:
            lines.append(f"✅ **{label}** — {w['model']} · llama.cpp, local")
        elif w["stage"] == "idle":
            lines.append(f"⚪ **{label}** — not loaded yet (auto-loads at startup)")
        else:
            lines.append(f"⏳ **{label}** — {w['stage']} ({w['ts']}): {w['detail'] or '…'}")
    lines.append("\n_Each sub-agent has its own lock — Explain / narrative / "
                 "research run in parallel without “model busy”._")
    return "\n\n".join(lines)


# ─────────────────────── inference ───────────────────────
DEFAULT_SYSTEM = ("You are a sub-agent of Chan Compass, a US-equity dashboard. "
                  "Answer in clear, concise English.")
MAX_PROMPT_CHARS = 3200


def _messages(user: str, system: str):
    return [{"role": "system", "content": system + " /no_think"},
            {"role": "user", "content": user[:MAX_PROMPT_CHARS]}]


def chat(user: str, max_tokens: int = 500, temperature: float = 0.3,
         system: str = DEFAULT_SYSTEM, worker: str = "interpreter") -> str:
    """Blocking chat on one sub-agent (used by pipeline/agent code)."""
    w = WORKERS[_wk(worker)]; worker = _wk(worker)
    if w["llm"] is None:
        return ""
    if not w["lock"].acquire(timeout=180):
        return f"({WORKER_LABEL[worker]} busy — try again in a moment)"
    try:
        out = w["llm"].create_chat_completion(
            messages=_messages(user, system),
            max_tokens=max_tokens, temperature=temperature)
        txt = out["choices"][0]["message"]["content"] or ""
        return _THINK_RE.sub("", txt).strip()
    except Exception as e:
        return f"(model error: {e})"
    finally:
        w["lock"].release()


def chat_stream(user: str, max_tokens: int = 500, temperature: float = 0.3,
                system: str = DEFAULT_SYSTEM, worker: str = "interpreter"):
    """Streaming chat on one sub-agent — yields cumulative text immediately."""
    w = WORKERS[_wk(worker)]; worker = _wk(worker)
    if w["llm"] is None:
        yield (f"⏳ {WORKER_LABEL[worker]} isn't ready yet — "
               f"stage: {w['stage']}. Check the **Model** tab.")
        return
    if not w["lock"].acquire(timeout=5):
        yield f"⏳ {WORKER_LABEL[worker]} is finishing another answer — try again in a few seconds."
        return
    try:
        acc = ""
        for chunk in w["llm"].create_chat_completion(
                messages=_messages(user, system),
                max_tokens=max_tokens, temperature=temperature, stream=True):
            delta = chunk["choices"][0]["delta"].get("content") or ""
            if not delta:
                continue
            acc += delta
            yield _THINK_RE.sub("", acc).replace("<think>", "").strip()
        if not acc.strip():
            yield "(model returned no text — try again)"
    except Exception as e:
        yield f"(model error: {e})"
    finally:
        w["lock"].release()


def quick_test() -> str:
    """Sanity check both sub-agents."""
    import time
    outs = []
    for key in ("interpreter", "narrator", "reporter", "analyst"):
        if WORKERS[key]["llm"] is None:
            outs.append(f"{WORKER_LABEL[key]}: not loaded ({WORKERS[key]['stage']})")
            continue
        t0 = time.time()
        out = chat("Reply with exactly: OK", max_tokens=6, temperature=0.0, worker=key)
        outs.append(f"{WORKER_LABEL[key]}: **{out or '(no output)'}** · {time.time()-t0:.1f}s")
    return "\n\n".join(outs)