#!/usr/bin/env python3
"""倍智器 · 自更新（Ubuntu；只用 Python 标准库，不依赖 curl/wget）

用法:
    beizhiqi-update check        # 只看有没有新版，不动任何文件
    beizhiqi-update update       # 有新版就更新（默认动作）
    beizhiqi-update rollback     # 退回到上一个版本
选项:
    --prefix DIR   安装根目录（默认 /opt/beizhiqi）
    --base URL     更新源（默认 https://dl.linkclaw.tech/beizhiqi）
    --force        版本号相同也重装一次

设计要点（Linux 上为什么可以很简单）:
  1. 版本各占一个目录 versions/<ver>/，靠软链 current 指向在用版本
  2. 切换 = 换软链（原子操作）；正在运行的进程用的是旧 inode，不受影响
  3. 任何一步失败都不动 current —— 更新失败不影响继续用旧版
  4. 下载后必须过 sha256，否则一律不装（防止下到半截或被人换包）
"""
import hashlib
import json
import os
import shutil
import subprocess
import sys
import tarfile
import tempfile
import urllib.request

DEFAULT_PREFIX = "/opt/beizhiqi"
DEFAULT_BASE = "https://dl.linkclaw.tech/beizhiqi"
LIST_TIMEOUT = 10      # 拉清单最多等 10 秒：拉不到就当没网，立刻返回
BIN_TIMEOUT = 30


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


def read_version(prefix, link="current"):
    p = os.path.join(prefix, link, "VERSION")
    if not os.path.isfile(p):
        return ""
    return open(p, encoding="utf-8").read().strip()


def fetch_manifest(base, timeout=LIST_TIMEOUT):
    url = base.rstrip("/") + "/latest.json"
    req = urllib.request.Request(url, headers={"User-Agent": "beizhiqi-update/1.0"})
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.loads(r.read().decode("utf-8"))


def newer(remote, local):
    # 用点分数字比较；解析不出来就按字符串比
    def key(v):
        try:
            return [int(x) for x in v.split(".")]
        except ValueError:
            return v
    try:
        return key(remote) > key(local)
    except TypeError:
        return remote != local


def download(url, dest, timeout=1800):
    req = urllib.request.Request(url, headers={"User-Agent": "beizhiqi-update/1.0"})
    with urllib.request.urlopen(req, timeout=timeout) as r, open(dest, "wb") as f:
        shutil.copyfileobj(r, f, length=1 << 20)


def extract(archive, dest):
    """解包并去掉单层顶层目录（发布包结构：<任意名>/bin/... VERSION）"""
    with tempfile.TemporaryDirectory() as tmp:
        with tarfile.open(archive, "r:gz") as t:
            t.extractall(tmp)
        entries = os.listdir(tmp)
        root = os.path.join(tmp, entries[0]) if len(entries) == 1 else tmp
        shutil.copytree(root, dest)


def smoke(bindir):
    exe = os.path.join(bindir, "beizhiqi")
    if not os.path.isfile(exe):
        return False, "包里没有 bin/beizhiqi"
    os.chmod(exe, 0o755)
    try:
        r = subprocess.run([exe, "--version"], capture_output=True, timeout=BIN_TIMEOUT)
    except Exception as e:                       # noqa: BLE001 — 冒烟测试失败都算失败
        return False, "跑不起来：%s" % e
    return r.returncode == 0, (r.stdout or b"").decode("utf-8", "replace").strip()


def switch(prefix, ver):
    """原子切换 current，并保留 previous 以便回退"""
    cur = os.path.join(prefix, "current")
    prev = os.path.join(prefix, "previous")
    if os.path.islink(cur) and os.path.exists(cur):
        if os.path.islink(prev):
            os.unlink(prev)
        os.symlink(os.readlink(cur), prev)
    tmp = os.path.join(prefix, ".current.new")
    if os.path.islink(tmp):
        os.unlink(tmp)
    os.symlink(os.path.join(prefix, "versions", ver), tmp)
    os.replace(tmp, cur)          # 原子替换


def main():
    args = [a for a in sys.argv[1:]]
    mode = "update"
    for a in args:
        if a in ("check", "update", "rollback"):
            mode = a
    prefix = DEFAULT_PREFIX
    base = DEFAULT_BASE
    force = "--force" in args
    if "--prefix" in args:
        prefix = args[args.index("--prefix") + 1]
    if "--base" in args:
        base = args[args.index("--base") + 1]

    local = read_version(prefix)
    if not local:
        print("没找到已安装的版本（%s/current/VERSION）——请先用 install.sh 安装" % prefix)
        return 1

    if mode == "rollback":
        prevp = os.path.join(prefix, "previous")
        if not os.path.islink(prevp):
            print("没有上一个版本可退")
            return 1
        ver = os.path.basename(os.readlink(prevp))
        switch(prefix, ver)
        print("已退回 %s（当前版本 %s）" % (ver, read_version(prefix)))
        return 0

    try:
        man = fetch_manifest(base)
    except Exception as e:                       # noqa: BLE001 — 没网就安静退出
        print("拉不到更新清单（%s）：%s" % (base, e))
        return 3 if mode == "check" else 0

    remote = str(man.get("version", ""))
    print("本机版本 %s · 服务器版本 %s" % (local, remote))
    if not remote or not (newer(remote, local) or (force and remote == local)):
        print("已是最新，不处理")
        return 0
    if mode == "check":
        print("有新版本可更新")
        return 0

    files = man.get("files", [])
    want = man.get("package", "")
    item = None
    if want:
        item = next((f for f in files if f.get("name") == want), None)
    if item is None:                      # 老清单没有 package 字段时按版本号挑
        item = next((f for f in files if f.get("name", "").endswith(".tar.gz") and remote in f["name"]), None)
    if item is None:
        item = next((f for f in files if f.get("name", "").endswith(".tar.gz")), None)
    if not item:
        print("清单里没有 .tar.gz 包，停止")
        return 1

    updir = os.path.join(prefix, "updates")
    os.makedirs(updir, exist_ok=True)
    pkg = os.path.join(updir, item["name"])
    print("下载 %s（%.1f MB）…" % (item["name"], item.get("size", 0) / 1048576.0))
    try:
        download(item["url"], pkg)
    except Exception as e:                       # noqa: BLE001
        print("下载失败：%s（保持当前版本不变）" % e)
        return 1

    got = sha256(pkg)
    if got != item.get("sha256"):
        os.unlink(pkg)
        print("校验不通过（应为 %s…，实际 %s…），已删除下载文件，保持当前版本"
              % (item["sha256"][:16], got[:16]))
        return 1

    target = os.path.join(prefix, "versions", remote)
    tmpdir = target + ".new"
    shutil.rmtree(tmpdir, ignore_errors=True)
    extract(pkg, tmpdir)
    ok, info = smoke(os.path.join(tmpdir, "bin"))
    if not ok:
        shutil.rmtree(tmpdir, ignore_errors=True)
        print("新版冒烟测试没过：%s（保持当前版本不变）" % info)
        return 1
    shutil.rmtree(target, ignore_errors=True)
    os.rename(tmpdir, target)
    switch(prefix, remote)
    os.chmod(pkg, 0o644)
    print("已更新到 %s（%s）" % (remote, info))
    return 0


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