#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""colab_fetch.py — Colab T4 弹药回拉: adapter+mmproj+底模(18G) -> 自动点火 colab_deploy.py
坑#57 对策: snapshot_download 走 cache_dir + symlink, 避免 local_dir 嵌套
幂等旗标: /content/k8b/FETCH_* ; 日志: /content/k8b/colab_fetch.log"""
import os, subprocess, time, traceback

K8B = "/content/k8b"
K27 = "/content/k27"
BASE = K27 + "/Qwen3.5-9B"
LOG = K8B + "/colab_fetch.log"
os.makedirs(K8B, exist_ok=True)
log = open(LOG, "a", buffering=1)

def P(m):
    log.write("[%s] %s\n" % (time.strftime("%m-%d %H:%M:%S"), m))

def sh(c, t=1800):
    p = subprocess.run(c, shell=True, capture_output=True, text=True, timeout=t, errors="replace")
    return ((p.stdout or "") + (p.stderr or ""))[-800:]

def flag(n, c="1"):
    open(K8B + "/" + n, "w").write(str(c)[:400])

def have(n):
    return os.path.exists(K8B + "/" + n)

def hf_token():
    return open("/root/.cache/huggingface/token").read().strip()

def main():
    t0 = time.time()
    P("==== colab fetch start ====")
    try:
        from huggingface_hub import hf_hub_download, snapshot_download, HfApi
        import shutil
        tok = hf_token()
        api = HfApi(token=tok)
        # F1: 脚本弹药 (公仓 artifacts/m15_scripts/)
        if not have("FETCH_SCRIPTS"):
            os.makedirs(K8B, exist_ok=True)
            for rp, dst in (
                ("artifacts/m15_scripts/colab_compile.py", K8B + "/colab_compile.py"),
                ("artifacts/m15_scripts/colab_deploy.py", K8B + "/colab_deploy.py"),
                ("artifacts/m15_scripts/r12_surgery.py", K8B + "/r12_surgery.py"),
            ):
                if not os.path.exists(dst):
                    p = hf_hub_download("tchbcb/samai-9b", rp, token=tok)
                    shutil.copy(p, dst)
                    P("got %s (%d)" % (dst, os.path.getsize(dst)))
            flag("FETCH_SCRIPTS")
        # F2: r10 adapter (私仓 242M)
        if not have("FETCH_ADP"):
            for f in ("adapter_config.json", "adapter_model.safetensors",
                      "chat_template.jinja", "tokenizer.json", "tokenizer_config.json"):
                d = K8B + "/m12_adapters/m13_r10"
                os.makedirs(d, exist_ok=True)
                if not os.path.exists(d + "/" + f):
                    p = hf_hub_download("tchbcb/samai-8b-M8", "m12_adapters/m13_r10/" + f, token=tok)
                    shutil.copy(p, d + "/" + f)
            n = len(os.listdir(K8B + "/m12_adapters/m13_r10"))
            if n < 5:
                raise RuntimeError("adapter incomplete: %d files" % n)
            P("adapter ok (%d files)" % n)
            flag("FETCH_ADP", str(n))
        # F3: mmproj (公仓 0.92G)
        if not have("FETCH_MM"):
            if not os.path.exists(K8B + "/mmproj_m11.gguf"):
                p = hf_hub_download("tchbcb/samai-9b", "mmproj_m11.gguf", token=tok)
                shutil.copy(p, K8B + "/mmproj_m11.gguf")
            P("mmproj ok (%.2fG)" % (os.path.getsize(K8B + "/mmproj_m11.gguf") / 1e9))
            flag("FETCH_MM")
        # F4: 底模 18G (cache_dir 隔离 + symlink 防嵌套)
        if not have("FETCH_BASE"):
            if not os.path.exists(BASE + "/config.json"):
                cache = K27 + "/hf"
                os.makedirs(cache, exist_ok=True)
                P("snapshot_download base (18G)...")
                p = snapshot_download("Qwen/Qwen3.5-9B", cache_dir=cache, token=tok)
                if os.path.islink(BASE):
                    os.remove(BASE)
                elif os.path.isdir(BASE):
                    shutil.rmtree(BASE)
                os.symlink(p, BASE)
                total = sum(os.path.getsize(os.path.join(r, f))
                            for r, _, fs in os.walk(p) for f in fs)
                P("base ok %s (%.1fG)" % (p, total / 1e9))
                if total < 15e9:
                    raise RuntimeError("base too small: %.1fG" % (total / 1e9))
            flag("FETCH_BASE")
        flag("FETCH_ALL_DONE", "%.0fs" % (time.time() - t0))
        P("==== FETCH_ALL_DONE, firing deploy ====")
        # F5: 自动点火部署链 (分离)
        sh("cd %s && rm -f DEPLOY_FAIL && setsid nohup python3 colab_deploy.py > /dev/null 2>&1 & sleep 1", 20)
        P("deploy fired")
    except Exception as e:
        traceback.print_exc(file=log)
        flag("FETCH_FAIL", repr(e)[:200])
        P("==== FETCH_FAIL: %s ====" % repr(e)[:200])

if __name__ == "__main__":
    main()