#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""倍智器 · 引擎按需下载器（渐进加载）

设计：安装包只带外壳，引擎不打包（几百 MB，且必须与显卡架构对上）。
装完之后按需拉、**支持断点续传**（分块下载，断了接着下，不重来）。

    beizhiqi-engine list                 # 看有哪些引擎、本机该装哪个
    beizhiqi-engine get <名字>            # 下载并安装到当前版本目录
    beizhiqi-engine get --auto            # 按本机显卡自动挑一个
    beizhiqi-engine verify                # 检查已装引擎是否完整（sha256）

只用 Python 标准库，不依赖 curl/wget。下载地址来自更新源清单 latest.json 的 components 段。
"""
import json, os, re, shutil, subprocess, sys, tarfile, urllib.request

BASE = os.environ.get("BEIZHIQI_BASE", "https://dl.linkclaw.tech/beizhiqi")
PREFIX = os.environ.get("BEIZHIQI_PREFIX", "/opt/beizhiqi")
CHUNK = 4 << 20          # 4 MB 一块，断点续传的最小粒度
TIMEOUT = 30


def manifest():
    url = BASE.rstrip("/") + "/latest.json"
    with urllib.request.urlopen(url, timeout=TIMEOUT) as r:
        return json.loads(r.read())


def components():
    return (manifest().get("components") or {}).get("engine", [])


def current_dir():
    cur = os.path.join(PREFIX, "current")
    if os.path.islink(cur):
        cur = os.path.realpath(cur)
    if not os.path.isdir(cur):
        sys.exit("没找到已安装的版本目录（%s），先跑 install.sh" % cur)
    return cur


def installed_engine():
    p = os.path.join(current_dir(), "bin", "llama-server")
    return p if os.path.isfile(p) else None


def gpu_hint():
    """本机显卡 → 该装哪个引擎。只看驱动报的算力/型号，够用就行。"""
    try:
        out = subprocess.run(["nvidia-smi", "--query-gpu=name,compute_cap", "--format=csv,noheader"],
                             capture_output=True, text=True, timeout=15).stdout.strip()
    except Exception:
        return None, "没检测到 NVIDIA 显卡（nvidia-smi 不可用）"
    if not out:
        return None, "没检测到 NVIDIA 显卡"
    caps = [l.split(",")[-1].strip() for l in out.splitlines()]
    for cap, tag in (("12", "sm120"), ("8.9", "sm86-89"), ("8.6", "sm120"), ("8.0", "sm86-89")):
        for c in caps:
            if c.startswith(cap):
                return tag, "本机显卡算力 %s" % ",".join(caps)
    return None, "显卡算力 %s 未匹配到已知引擎" % ",".join(caps)


def pick(comps, tag):
    for c in comps:
        if tag and tag in c["name"]:
            return c
    return None


def download(url, dest, size=None):
    """分块下载 + 断点续传：已下多少就从哪继续（服务端支持 Range 才行，nginx 支持）"""
    have = os.path.getsize(dest) if os.path.exists(dest) else 0
    req = urllib.request.Request(url)
    if have:
        req.add_header("Range", "bytes=%d-" % have)
    with urllib.request.urlopen(req, timeout=TIMEOUT) as r:
        if have and r.status != 206:          # 服务端不给续传，重来
            have = 0
            os.remove(dest)
        total = size or (have + int(r.headers.get("Content-Length") or 0))
        mode = "ab" if have else "wb"
        done = have
        with open(dest, mode) as f:
            while True:
                buf = r.read(CHUNK)
                if not buf:
                    break
                f.write(buf)
                done += len(buf)
                print("\r  已下载 %.1f / %.1f MB" % (done / 1048576, total / 1048576), end="", flush=True)
    print()
    return os.path.getsize(dest)


def sha256(path):
    import hashlib
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for b in iter(lambda: f.read(1 << 20), b""):
            h.update(b)
    return h.hexdigest()


def cmd_list():
    comps = components()
    if not comps:
        print("更新源里还没有引擎条目")
        return
    tag, why = gpu_hint()
    print("可用引擎（%s）：" % why)
    for c in comps:
        mark = "  ← 本机建议" if tag and tag in c["name"] else ""
        print("  %-34s %6.1f MB  %s%s" % (c["name"], c["size"] / 1048576, c.get("note", ""), mark))
    ins = installed_engine()
    print("本机已装引擎：%s" % (ins or "无（现在只有外壳，还不能跑模型）"))


def cmd_get(name=None, auto=False):
    comps = components()
    if not comps:
        sys.exit("更新源里还没有引擎条目")
    c = None
    if auto:
        tag, why = gpu_hint()
        c = pick(comps, tag)
        print(why)
        if not c:
            sys.exit("没有匹配本机显卡的引擎，用 `beizhiqi-engine list` 看全部，再手动指定")
    else:
        c = next((x for x in comps if x["name"] == name), None)
        if not c:
            sys.exit("没有这个引擎：%s（跑 beizhiqi-engine list 看可选）" % name)
    tmp = os.path.join(PREFIX, "updates", c["file"])
    os.makedirs(os.path.dirname(tmp), exist_ok=True)
    print("下载 %s（%.1f MB，可断点续传）" % (c["file"], c["size"] / 1048576))
    download(c["url"], tmp, c["size"])
    got = sha256(tmp)
    if got != c["sha256"]:
        sys.exit("sha256 对不上（期望 %s，实得 %s）——文件不完整或被改过，重跑一次会自动续传" % (c["sha256"][:12], got[:12]))
    print("校验通过")
    dest = os.path.join(current_dir(), "bin")
    os.makedirs(dest, exist_ok=True)
    with tarfile.open(tmp, "r:gz") as t:
        t.extractall(dest)                     # 包里是 llama-server 等文件
    # 兼容包内多一层目录
    for root, _dirs, files in os.walk(dest):
        for f in files:
            src = os.path.join(root, f)
            tgt = os.path.join(dest, f)
            if src != tgt and not os.path.exists(tgt):
                os.replace(src, tgt)
    for f in os.listdir(dest):
        p = os.path.join(dest, f)
        if os.path.isfile(p):
            os.chmod(p, 0o755)
    print("引擎已装到 %s" % dest)
    subprocess.run([os.path.join(dest, "llama-server"), "--version"], check=False)


def cmd_verify():
    ins = installed_engine()
    if not ins:
        sys.exit("本机还没装引擎")
    print("已装引擎：%s（%.1f MB）" % (ins, os.path.getsize(ins) / 1048576))
    print(subprocess.run([ins, "--version"], capture_output=True, text=True).stdout.strip()[:200])


def demo():
    """自检：断点续传的分块逻辑 —— 起本地 http 服务，下 3 块后中断，再续完，比对 sha256"""
    import http.server, socketserver, tempfile, threading, hashlib
    d = tempfile.mkdtemp()
    payload = os.urandom(3 * 1024 * 1024 + 1234)
    src = os.path.join(d, "engine-t.tar.gz")
    open(src, "wb").write(payload)
    h = hashlib.sha256(payload).hexdigest()
    os.chdir(d)

    class H(http.server.SimpleHTTPRequestHandler):
        def log_message(self, *a):
            pass

    srv = socketserver.TCPServer(("127.0.0.1", 0), H)
    threading.Thread(target=srv.serve_forever, daemon=True).start()
    url = "http://127.0.0.1:%d/engine-t.tar.gz" % srv.server_address[1]

    dst = os.path.join(d, "part.bin")
    global CHUNK
    keep, CHUNK = CHUNK, 1024 * 1024
    download(url, dst, len(payload) * 2)      # 第一轮：只下 1 MB 就"断"
    with open(dst, "r+b") as f:               # 模拟中断：截断到 1MB
        f.truncate(1024 * 1024)
    download(url, dst)                        # 第二轮：应从头续传（Range）
    CHUNK = keep
    assert os.path.getsize(dst) == len(payload), "续传后长度不对"
    assert sha256(dst) == h, "续传后内容不对"
    srv.shutdown()
    print("自检通过（断点续传）")


if __name__ == "__main__":
    a = sys.argv[1:]
    if not a or a[0] in ("list", "ls"):
        cmd_list()
    elif a[0] == "get":
        if "--auto" in a:
            cmd_get(auto=True)
        elif len(a) > 1:
            cmd_get(name=a[1])
        else:
            print(__doc__)
    elif a[0] == "verify":
        cmd_verify()
    elif a[0] == "demo":
        demo()
    else:
        print(__doc__)
