#!/usr/bin/env python3
"""EPM 單一入口 demo：從「什麼都沒有」到「可以聊天」。

    python demo.py

會依序做這些事（每一步都可以重跑而不會壞掉：續傳、增量編譯、已存在就跳過）：

  1. 下載本專案的 code（從 Hugging Face model repo）
     —— 如果你已經在本專案的原始碼目錄裡執行，就直接用現有的，不重下。
  2. 下載模型 `openbmb/MiniCPM-V-4.7-35B-A3B-gguf/MiniCPM-V-4_7-35B-A3B-Q4_K_M.gguf`
     （19.71 GB，支援續傳；中斷再跑會從斷點繼續）
  3. 取得 llama.cpp（**固定 commit**）+ 套用 EPM patch + 編譯
  4. 以 8 GiB 實體 RAM 預算啟動 `epq-chat`，進入互動式聊天

EPM pager budget 是「實體 RAM 預算」，不是虛擬記憶體上限，也不會停用系統 swap。

注意：SDQ_RESERVE_MB 是逐模型、逐 ctx/ubatch 校正出來的預留量。本模型在
ctx=4096 / ubatch=128 / 8 GiB 預算下是 512 MiB；換 RAM 預算必須重新校正
（見 docs/DESIGN.md §3）。
要取得嚴格 8 GiB / no-swap 的證據，請用 `tools/strict8g.py`（或
`python -m sdq.strict_bench`）把整個行程管治起來；方法與限制見
docs/STRICT_8GIB.md。

設計原則：任何一步都可以單獨重跑；失敗時不會留下壞掉的半成品。
"""

from __future__ import annotations

import argparse
import os
import shutil
import subprocess
import sys
import time
from pathlib import Path

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))

try:
    from sdq import config as C
except Exception:  # 允許在 code 尚未下載完時也能顯示說明
    C = None

DEFAULT_ROOT = Path(os.environ.get(
    "EPM_HOME", os.environ.get("SDQ_HOME", str(Path.home() / "minicpm47-epm"))
))

# 從 HF 拉回來的檔案（給 fetch_code 用）
# 從 HF 拉回來的檔案（給 fetch_code 用）。
# tools/*.c 也要拉：tools/strict8g.py 會在執行時把它編譯成 LD_PRELOAD helper。
CODE_PATTERNS = [
    "*.py", "*.md", "*.c",
    "patches/*", "cpp/*", "docs/*", "tools/*", "sdq/*.py",
    "bench/expert_index.json",
]


def log(msg: str) -> None:
    print(f"[demo] {msg}", flush=True)


def run(cmd: list, **kw) -> subprocess.CompletedProcess:
    log("$ " + " ".join(str(c) for c in cmd))
    return subprocess.run([str(c) for c in cmd], **kw)


def human(n: float) -> str:
    for u in ("B", "KiB", "MiB", "GiB", "TiB"):
        if abs(n) < 1024:
            return f"{n:.2f} {u}"
        n /= 1024
    return f"{n:.1f} PiB"


def cfg(name: str, fallback):
    return getattr(C, name, fallback) if C else fallback


def reload_config(code_dir: Path) -> None:
    """code 是從 HF 下載到別的目錄時，改用那份 code 裡的 sdq/config.py。"""
    global C
    if str(code_dir) not in sys.path:
        sys.path.insert(0, str(code_dir))
    for mod in [m for m in list(sys.modules) if m == "sdq" or m.startswith("sdq.")]:
        del sys.modules[mod]
    try:
        from sdq import config as new_c
        C = new_c
    except Exception:
        pass


# ---------------------------------------------------------------- 1) 取得 code
def fetch_code(dest: Path, force: bool = False) -> Path:
    """如果目前就在本專案的原始碼裡，直接回傳；否則從 HF 下載。"""
    if (dest / "cpp" / "sdq_pager.cpp").exists() and (dest / cfg("PATCH_FILE", "")).exists() \
            and not force:
        log(f"已經在 EPM 專案原始碼中：{dest}")
        return dest
    log(f"從 Hugging Face 下載專案 code：{cfg('HF_REPO', 'HelloSun/minicpm47woek')}")
    from huggingface_hub import snapshot_download

    dest.mkdir(parents=True, exist_ok=True)
    snapshot_download(
        repo_id=cfg("HF_REPO", "HelloSun/minicpm47woek"),
        repo_type="model",
        local_dir=dest,
        allow_patterns=CODE_PATTERNS,
        token=os.environ.get("HF_TOKEN"),
    )
    if not (dest / cfg("PATCH_FILE", "")).exists():
        raise SystemExit(
            f"下載完成但缺少 {cfg('PATCH_FILE', '')}；請確認 HF 上的 code 完整，"
            "或直接在專案原始碼目錄執行 demo.py"
        )
    log(f"code 已下載到 {dest}")
    return dest


# ---------------------------------------------------------------- 2) 模型
def fetch_model(models_dir: Path, force: bool = False) -> Path:
    repo = cfg("MODEL_REPO", "openbmb/MiniCPM-V-4.7-35B-A3B-gguf")
    name = cfg("MODEL_FILE", "MiniCPM-V-4_7-35B-A3B-Q4_K_M.gguf")
    size = cfg("MODEL_SIZE", 21_166_255_008)
    path = models_dir / name
    if path.exists() and not force and path.stat().st_size == size:
        log(f"模型已存在：{path}（{human(path.stat().st_size)}）")
        return path

    log(f"下載模型 {repo}/{name}（{human(size)}，可續傳）")
    log("    需要約 22 GB 磁碟空間；下載到一半再執行本程式會從斷點繼續")
    from huggingface_hub import hf_hub_download

    models_dir.mkdir(parents=True, exist_ok=True)
    p = hf_hub_download(repo, name, local_dir=str(models_dir),
                        token=os.environ.get("HF_TOKEN"))
    got = Path(p)
    if got.stat().st_size != size:
        log(f"⚠️ 檔案大小 {got.stat().st_size} 與預期 {size} 不同，請再跑一次本程式續傳")
    log(f"模型完成：{got}")
    return got


# ---------------------------------------------------------------- 3) 原始碼 + 編譯
def fetch_source(src_dir: Path, force: bool = False) -> Path:
    """取得固定 commit 的 llama.cpp。已存在且 commit 正確就跳過。"""
    commit = cfg("LLAMA_COMMIT", "836d57176dc699a726c55418e4f96b8ca628e1bf")
    repo = cfg("LLAMA_REPO", "https://github.com/ggml-org/llama.cpp.git")
    src_dir.mkdir(parents=True, exist_ok=True)

    if (src_dir / ".git").exists() and not force:
        cur = subprocess.run(["git", "rev-parse", "HEAD"], cwd=src_dir,
                             capture_output=True, text=True).stdout.strip()
        if cur == commit:
            log(f"llama.cpp 原始碼已存在且版本正確：{src_dir}")
            return src_dir
        log(f"版本不符（目前 {cur[:10]}），重新取得 {commit[:10]}")
        shutil.rmtree(src_dir / ".git")
        shutil.rmtree(src_dir / "src", ignore_errors=True)

    if not (src_dir / ".git").exists():
        log(f"clone llama.cpp（固定 commit {commit[:10]}）")
        tmp = src_dir.parent / (src_dir.name + ".tmp")
        shutil.rmtree(tmp, ignore_errors=True)
        run(["git", "clone", "--filter=blob:none", "--no-checkout", repo, tmp],
            check=True)
        # GitHub 允許依 commit 直接 fetch，失敗就退回淺層 clone 再 checkout
        r = subprocess.run(["git", "fetch", "-q", "--depth", "1", "origin", commit],
                           cwd=tmp, capture_output=True, text=True)
        if r.returncode != 0:
            log("依 commit 直接 fetch 失敗，改用完整 clone")
            run(["git", "fetch", "-q", "--tags", "origin"], cwd=tmp, check=True)
            run(["git", "fetch", "-q", "origin"], cwd=tmp, check=True)
        run(["git", "checkout", "-q", commit], cwd=tmp, check=True)
        shutil.rmtree(src_dir, ignore_errors=True)
        tmp.rename(src_dir)
    return src_dir


def apply_patch(src_dir: Path, code_dir: Path = None) -> None:
    """把 SDQ 的 expert pager 引擎套進 llama.cpp 的 src/。

    `code_dir` 是本專案 code 所在的位置（可能不是 demo.py 旁邊，例如 code 是從 HF
    下載到 <home>/code 的情況）。
    """
    rel = cfg("PATCH_FILE", "patches/0001-sdq-expert-pager.patch")
    patch = (code_dir or HERE) / rel
    target = src_dir / cfg("PATCH_TARGET_DIR", "src")
    marker = src_dir / ".sdq_patched"

    if marker.exists() and all((target / f).exists()
                               for f in cfg("SDQ_SOURCE_FILES", [])):
        log("SDQ patch 已套用過")
        return
    if not patch.exists():
        raise SystemExit(f"找不到 patch：{patch}")
    if not all((target / f).exists() for f in cfg("SDQ_SOURCE_FILES", [])):
        log("套用 SDQ expert pager patch（引擎檔案會進到 llama.cpp/src/）")
    r = subprocess.run(["git", "apply", "--check", str(patch)], cwd=src_dir,
                       capture_output=True, text=True)
    if r.returncode != 0:
        # patch 已經在樹裡（常見於換機／換目錄重建）
        if all((target / f).exists() for f in cfg("SDQ_SOURCE_FILES", [])):
            log("patch 已在樹中（略過重複套用）")
            marker.touch()
            return
        raise SystemExit(f"patch 套用失敗：{r.stderr.strip()}")
    subprocess.run(["git", "apply", str(patch)], cwd=src_dir, check=True)
    # Keep the checked-in pager implementation in sync with the patch's
    # original source payload. This also lets incremental rebuilds consume
    # optional backends added after the historical patch was generated.
    for name in cfg("SDQ_SOURCE_FILES", []):
        source = (code_dir or HERE) / "cpp" / name
        target_file = target / name
        if source.exists() and target_file.exists():
            shutil.copy2(source, target_file)
    marker.touch()
    log("patch 套用完成")


def build(src_dir: Path, jobs: int, force: bool = False) -> Path:
    target = cfg("BUILD_TARGET", "epq-chat")
    binary = src_dir / "build" / "bin" / target
    if binary.exists() and not force:
        log(f"已經編譯過：{binary}")
        return binary

    build_dir = src_dir / "build"
    log("設定 cmake（CPU 最佳化：AVX2/AVX512 由 GGML_NATIVE 自動開啟）")
    run(["cmake", "-B", str(build_dir), "-S", str(src_dir),
         "-DCMAKE_BUILD_TYPE=Release", "-DGGML_NATIVE=ON",
         "-DLLAMA_CURL=OFF", "-DLLAMA_BUILD_TESTS=OFF",
         "-DLLAMA_BUILD_EXAMPLES=OFF", "-DLLAMA_BUILD_SERVER=OFF"],
        check=True)
    log(f"編譯中（-j{jobs}），參考實測約 2-5 分鐘")
    t0 = time.time()
    run(["cmake", "--build", str(build_dir), "-j", str(jobs), "--target", target],
        check=True)
    log(f"編譯完成（{time.time() - t0:.0f}s）：{binary}")
    return binary


def reserve_mb(ram_mb) -> int:
    """arena 之外要預留多少。

    預留量必須涵蓋 arena sizing 之後才長出來的記憶體：非專家權重第一次 forward
    才 fault in、CPU compute buffer、KV cache、圖與執行緒堆疊。這是逐模型、
    逐 ctx/ubatch 校正出來的量，這裡用一條比例規則讓不同 RAM 預算都有合理餘量。
    """
    ram = int(ram_mb)
    if C is not None and getattr(C, "RESERVE_MB", None) and ram == int(C.RAM_BUDGET_MB):
        return int(C.RESERVE_MB)
    return max(128, min(1024, ram // 16))


# ---------------------------------------------------------------- 4) 執行
def chat(binary: Path, model: Path, args) -> int:
    cmd = [
        binary, "-m", model,
        "-c", args.ctx, "-t", args.threads, "-n", args.n_predict,
        "--ram-budget-mb", args.ram_mb,
    ]
    if args.arena_mb:
        cmd += ["--arena-mb", args.arena_mb]
    if args.no_prefetch:
        cmd += ["--no-prefetch"]
    if args.stats:
        cmd += ["--stats-file", args.stats]
    # SDQ_RESERVE_MB 是逐模型、逐 ctx/ubatch 校正出來的預留量，預設沒有設定。
    # 這裡依 config 的比例規則帶入，否則使用者拿到的是另一個模型的校正值。
    reserve = reserve_mb(args.ram_mb)
    env = dict(os.environ, SDQ_RESERVE_MB=str(reserve))
    log(f"啟動 epq-chat（RAM 預算 {args.ram_mb} MiB，reserve {reserve} MiB；"
        f"改 ctx 或 ubatch 必須重新校正，見 docs/DESIGN.md §3）")
    return subprocess.run([str(c) for c in cmd], env=env).returncode


def main() -> int:
    ap = argparse.ArgumentParser(
        description="MiniCPM-V-4.7-35B-A3B EPM CPU inference: download → build → chat",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""常用範例：
  python demo.py                    # 完整流程，然後進入聊天
  python demo.py --bench            # 不聊天，直接跑基準測量（量 RSS / SSD I/O / 命中率 / tok/s）
  python demo.py --verify           # 驗證分頁路徑的數值正確性（A/B + NumPy）
  python demo.py --verify --fast    # 只跑 8 GiB 內做得完的層（略過端到端 A/B）
  python demo.py --once             # 不聊天，單跑一句示範（適合 CI / 腳本）
  python demo.py --model-path /x.gguf    # 模型已經在自己電腦上
  python demo.py --rebuild          # 強制重新編譯 llama.cpp
""")
    ap.add_argument("--home", default=str(DEFAULT_ROOT), help="工作目錄（預設 ~/minicpm47-epm）")
    ap.add_argument("--src", help="llama.cpp 原始碼目錄（預設 <home>/llama.cpp）")
    ap.add_argument("--model-path", help="已下載好的模型路徑")
    ap.add_argument("--ctx", default=str(cfg("CTX_VERIFY", 4096)),
                    help="上下文長度（模型原生 262144，這裡實測 4K）")
    ap.add_argument("--threads", type=int, default=0, help="CPU 執行緒（0 = 自動）")
    ap.add_argument("--n-predict", type=int, default=cfg("N_PREDICT", 256),
                    help="每次回答產生多少 token")
    ap.add_argument("--ram-mb", default=str(cfg("RAM_BUDGET_MB", 8192)),
                    help="SDQ pager sizing budget in MiB (not a total RSS/swap limit)")
    ap.add_argument("--arena-mb", type=int, default=cfg("ARENA_MB", 0),
                    help="固定 arena 大小（0 = 由 RAM 預算推算）")
    ap.add_argument("--jobs", type=int, default=max(1, (os.cpu_count() or 4)), help="編譯並行度")
    ap.add_argument("--no-prefetch", action="store_true", help="關閉背景預取")
    ap.add_argument("--skip-code", action="store_true", help="不要抓 HF 上的 code")
    ap.add_argument("--skip-model", action="store_true", help="不要下載模型")
    ap.add_argument("--skip-build", action="store_true", help="不要編譯")
    ap.add_argument("--rebuild", action="store_true", help="強制重新編譯")
    ap.add_argument("--bench", action="store_true", help="跑完就做基準測量，不進聊天")
    ap.add_argument("--verify", action="store_true", help="做正確性驗證（A/B 對照 + NumPy 逐層比對）")
    ap.add_argument("--fast", action="store_true",
                    help="--verify 時只跑 8 GiB 內做得完的層：略過端到端 A/B"
                         "（那一層需要整個 22 GiB 模型常駐記憶體）")
    ap.add_argument("--once", action="store_true", help="不聊天，單跑一句示範就結束")
    ap.add_argument("--stats", default="", help="把指標寫到指定 JSONL")
    args = ap.parse_args()

    home = Path(args.home).expanduser().resolve()
    home.mkdir(parents=True, exist_ok=True)

    # --- 1) code：已在專案原始碼裡就直接用，否則從 HF 抓
    in_source_tree = (HERE / "cpp" / "sdq_pager.cpp").exists()
    if in_source_tree and not args.skip_code:
        code_dir = HERE
        log(f"code：使用目前的專案原始碼 {code_dir}")
    elif args.skip_code:
        code_dir = HERE
    else:
        code_dir = fetch_code(home / "code")

    log(f"工作目錄：{home}")
    log(f"專案 code：{code_dir}")
    reload_config(code_dir)

    # --- 2) 模型
    model = Path(args.model_path) if args.model_path else (home / "models" / cfg("MODEL_FILE", ""))
    if not args.skip_model and not args.model_path:
        model = fetch_model(model.parent)

    # --- 3) 原始碼 + patch + 編譯
    src_dir = Path(args.src) if args.src else (home / "llama.cpp")
    target = cfg("BUILD_TARGET", "epq-chat")
    if args.skip_build:
        binary = src_dir / "build" / "bin" / target
    else:
        fetch_source(src_dir, force=args.rebuild)
        apply_patch(src_dir, code_dir)
        binary = build(src_dir, args.jobs, force=args.rebuild)

    if not binary.exists():
        raise SystemExit(f"找不到執行檔：{binary}")

    sys.path.insert(0, str(code_dir))

    # --- 4a) 正確性驗證
    if args.verify:
        from sdq.verify import run_verification
        if args.fast:
            log("--fast：只跑 kernel / 整層 / 預取壓力；跳過端到端 A/B"
                "（那一層需要整個 22 GiB 模型常駐 RAM，8 GiB 機器上不可能）")
        return run_verification(binary, model, code_dir, threads=args.threads or 8,
                                fast=args.fast)

    # --- 4b) 基準測量
    if args.bench:
        from sdq.bench import main as bench_main
        sys.argv = ["sdq.bench", "--model", str(model), "--bin", str(binary),
                    "--ram-mb", args.ram_mb, "--ctx", args.ctx,
                    "--threads", str(args.threads or 8), "--label", "demo"]
        return bench_main()

    # --- 4c) 非互動模式（CI / 腳本）
    if args.once or not sys.stdin.isatty():
        if not args.once:
            log("沒有互動式終端，改用預設提示詞跑一次示範（加 --once 可明確指定）")
        cmd = [binary, "-m", model, "-c", args.ctx, "-t", args.threads,
               "--ram-budget-mb", args.ram_mb, "--no-chat", "--no-interactive",
               "--temp", "0", "-p", "請用一句話介紹你自己。", "-n", "64"]
        reserve = reserve_mb(args.ram_mb)
        log(f"單次示範：RAM 預算 {args.ram_mb} MiB，reserve {reserve} MiB"
            f"（改 ctx 或 ubatch 必須重新校正，見 docs/DESIGN.md §3）")
        log("注意：這個示範不受 strict8g 管治，實際 RSS 可能略高於 RAM 預算；"
            "受管治的量測請用 tools/strict8g.py")
        return subprocess.run([str(c) for c in cmd],
                              env=dict(os.environ, SDQ_RESERVE_MB=str(reserve))).returncode

    # --- 4d) 互動式聊天
    log("")
    log("=" * 68)
    log("開始聊天！模型 19.71 GB 只住在 SSD，RAM 只放有限容量的 expert 快取。")
    log("指令：/exit 離開　/reset 重設對話　/stats 看分頁指標（命中率 / SSD 讀取量）")
    log("=" * 68)
    return chat(binary, model, args)


if __name__ == "__main__":
    sys.exit(main())