#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""colab_compile.py — Colab T4 专用 llama.cpp CUDA 编译链 (sm_75)
编译 -> 打包 llama_cpp_cuda_colab_t4.tar.gz -> 第一时间推仓 (tchbcb/samai-9b)
幂等旗标: /content/k8b/CC_*
日志: /content/k8b/colab_compile.log"""
import os, subprocess, sys, time, traceback

K8B = "/content/k8b"
BIN = "/content/lcbin"
LCSRC = "/content/lcbuild/llama.cpp"
LOG = K8B + "/colab_compile.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=7200):
    p = subprocess.run(c, shell=True, capture_output=True, text=True, timeout=t, errors="replace")
    return ((p.stdout or "") + (p.stderr or ""))[-1500:]

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 compile start ====")
    try:
        # CC1: clone
        if not have("CC_CLONE"):
            if not os.path.exists(LCSRC + "/CMakeLists.txt"):
                r = sh("rm -rf /content/lcbuild && mkdir -p /content/lcbuild && "
                       "git clone --depth 1 https://github.com/ggml-org/llama.cpp %s" % LCSRC, 900)
                if not os.path.exists(LCSRC + "/CMakeLists.txt"):
                    raise RuntimeError("clone failed: " + r[-300:])
            flag("CC_CLONE", sh("cd %s && git rev-parse --short HEAD" % LCSRC, 15))
            P("clone ok")
        # CC2: cmake + build (sm_75, 2 核 -j2)
        if not have("CC_BUILD"):
            r = sh("cd /content/lcbuild && cmake -S %s -B %s/build -DGGML_CUDA=ON "
                   "-DCMAKE_CUDA_ARCHITECTURES=75 -DLLAMA_CURL=OFF -DCMAKE_BUILD_TYPE=Release "
                   "> %s/cmake.log 2>&1; echo RC=$?; tail -3 %s/cmake.log" % (LCSRC, LCSRC, LCSRC, LCSRC), 600)
            P("cmake: " + r[-200:])
            if "RC=0" not in r:
                raise RuntimeError("cmake failed")
            r = sh("cd %s/build && cmake --build . --config Release -j2 > %s/build.log 2>&1; echo RC=$?" % (LCSRC, LCSRC), 7200)
            P("build: " + r[-200:])
            if "RC=0" not in r:
                raise RuntimeError("build failed: " + sh("tail -20 %s/build.log" % LCSRC)[-600:])
            flag("CC_BUILD")
            P("build ok (%.0fs)" % (time.time() - t0))
        # CC3: 平铺 bin+so -> /content/lcbin
        if not have("CC_FLAT"):
            os.makedirs(BIN, exist_ok=True)
            sh("cp -a %s/build/bin/* %s/ 2>/dev/null" % (LCSRC, BIN), 300)
            sh("find %s/build -name '*.so*' -exec cp -a {} %s/ \\; 2>/dev/null" % (LCSRC, BIN), 300)
            need = ["llama-server", "llama-quantize", "llama-cli"]
            miss = [b for b in need if not os.path.exists(BIN + "/" + b)]
            if miss:
                raise RuntimeError("binaries missing: %s" % miss)
            r = sh("ls -la %s | head -20" % BIN, 15)
            P("flat ok:\n" + r)
            flag("CC_FLAT")
        # CC4: 打包 + 第一时间推仓
        if not have("CC_PUSH"):
            tar = K8B + "/llama_cpp_cuda_colab_t4.tar.gz"
            sh("tar czf %s -C %s . > /dev/null 2>&1" % (tar, BIN), 600)
            if os.path.getsize(tar) < 3e7:
                raise RuntimeError("tar too small: %d" % os.path.getsize(tar))
            from huggingface_hub import HfApi
            api = HfApi(token=hf_token())
            api.upload_file(path_or_fileobj=tar, path_in_repo="llama_cpp_cuda_colab_t4.tar.gz",
                            repo_id="tchbcb/samai-9b", repo_type="model",
                            commit_message="colab T4 (sm_75) CUDA build, first-push")
            flag("CC_PUSH", "%.0fs" % (time.time() - t0))
            P("CUDA tar pushed (%.0f MB, %.0fs)" % (os.path.getsize(tar) / 1e6, time.time() - t0))
        P("==== COMPILE_ALL_DONE in %.0fs ====" % (time.time() - t0))
        flag("COMPILE_ALL_DONE", "%.0fs" % (time.time() - t0))
    except Exception as e:
        traceback.print_exc(file=log)
        flag("COMPILE_FAIL", repr(e)[:200])
        P("==== COMPILE_FAIL: %s ====" % repr(e)[:200])

if __name__ == "__main__":
    main()