#!/usr/bin/env python3
"""Qoder UMID 设备风控身份（runtime-info）管理器。

背景：Qoder 服务端只向携带 Cosy-MachineToken / Cosy-MachineCode /
Cosy-MachineType 的请求下发「每天领 100 Credits」活动（社区实测，国际版尤其严格）。
这三个值由 Qoder 官方 UMID 程序 runtime-info 生成——桌面客户端自带，
官方 CLI 包 @qoder-ai/qodercli 也把各平台二进制以 base64 内嵌在 bundle/qodercli.js 中。

本模块在 NAS（Linux x86_64 / arm64）上下载官方 npm 包 → 提取对应架构的 ELF →
试跑一次确认可产出身份 → 缓存到应用数据目录（umid/）。

调用方式与客户端一致：
    runtime-info <env> [--account-stdin]，stdin 传入 {"account": "<uid>"}
    env: 国际版 4（SINGAPORE）/ 国内版 0（ONLINE）
第一行输出 JSON：{"machineToken","machineCode","machineType","vmInfo"...}

安全性：只从官方 npm registry 下载官方包，二进制仅本机使用，不上传任何信息。
"""
import base64
import json
import os
import re
import stat
import subprocess
import sys
import time
import urllib.request
import ssl

from common import DATA_DIR, log, USER_AGENT

PKG = "@qoder-ai/qodercli"
REGISTRIES = [
    "https://registry.npmmirror.com",
    "https://registry.npmjs.org",
]
BUNDLE_ENTRY = "package/bundle/qodercli.js"

# ELF e_machine: x86_64=62, arm64=183
ELF_MACHINE = {"x86_64": 62, "aarch64": 183, "arm64": 183}

RISK_ENV = {"qoder": 4, "qoder-cn": 0}   # 国际版 / 国内版

_SSL_CTX = ssl.create_default_context()
_SSL_CTX.check_hostname = False
_SSL_CTX.verify_mode = ssl.CERT_NONE

CACHE_DIR = os.path.join(DATA_DIR, "umid")
BIN_PATH = os.path.join(CACHE_DIR, "runtime-info")
MANIFEST = os.path.join(CACHE_DIR, "manifest.json")


# ---------------------------------------------------------------------------
# 缓存
# ---------------------------------------------------------------------------

def installed():
    """已安装的 runtime-info：{path, version, arch, sha256, installedAt} 或 None."""
    try:
        m = json.load(open(MANIFEST, encoding="utf-8"))
        if not os.path.isfile(BIN_PATH) or m.get("arch") != _arch():
            return None
        return dict(m, path=BIN_PATH)
    except (OSError, ValueError):
        return None


def _arch():
    machine = os.uname().machine
    return {"x86_64": "x86_64", "amd64": "x86_64",
            "aarch64": "aarch64", "arm64": "aarch64"}.get(machine, machine)


# ---------------------------------------------------------------------------
# 下载 / 提取
# ---------------------------------------------------------------------------

def _fetch(url, timeout=180, binary=True):
    req = urllib.request.Request(url, headers={"User-Agent": USER_AGENT})
    with urllib.request.urlopen(req, timeout=timeout, context=_SSL_CTX) as r:
        return r.read()


def _registry_meta(registry):
    meta = json.loads(_fetch(registry + "/" + PKG.replace("/", "%2F") + "/latest", 30).decode())
    return meta


def _extract_tar_entry(tar, wanted):
    """最小 tar 解析：按名字取一个普通文件内容。"""
    off = 0
    long_name = None
    while off + 512 <= len(tar):
        h = tar[off:off + 512]
        if not any(h):
            break
        name = h[:100].split(b"\0")[0].decode("utf-8", "replace")
        size = int((h[124:136].split(b"\0")[0].strip() or b"0"), 8)
        typeflag = h[156:157]
        prefix = h[345:500].split(b"\0")[0].decode("utf-8", "replace")
        full = (prefix + "/" + name) if prefix else name
        if typeflag == b"L":  # GNU long name
            long_name = tar[off + 512:off + 512 + size].split(b"\0")[0].decode("utf-8", "replace")
        elif typeflag in (b"0", b"\0"):
            real = long_name or full
            long_name = None
            if real == wanted:
                return tar[off + 512:off + 512 + size]
        off += 512 + ((size + 511) // 512) * 512
    return None


def _find_elf(bundle, arch):
    """在 bundle 里找指定架构的 ELF（base64 内嵌）。"""
    want = ELF_MACHINE.get(arch)
    if not want:
        raise RuntimeError("不支持的 CPU 架构：%s" % arch)
    for m in re.finditer(rb"[A-Za-z0-9+/]{100000,}", bundle):
        try:
            raw = base64.b64decode(m.group(0))
        except Exception:
            continue
        if raw[:4] != b"\x7fELF" or len(raw) < 18:
            continue
        machine = int.from_bytes(raw[18:20], "little")
        if machine == want:
            return raw
    raise RuntimeError("官方包中未找到 %s 架构的 UMID 二进制" % arch)


def install(force=False):
    """下载官方 qodercli 包并提取 runtime-info。返回 (ok, message)."""
    if not force and installed():
        return True, "已安装"
    os.makedirs(CACHE_DIR, exist_ok=True)
    arch = _arch()
    if arch not in ELF_MACHINE:
        return False, "NAS CPU 架构 %s 暂不支持（仅 x86_64 / arm64）" % arch
    last_err = None
    for registry in REGISTRIES:
        try:
            log("umid: fetching %s from %s" % (PKG, registry), "INFO")
            meta = _registry_meta(registry)
            version = meta.get("version", "")
            tarball = (meta.get("dist") or {}).get("tarball") or ""
            if not tarball:
                raise RuntimeError("registry 未返回 tarball")
            log("umid: downloading %s %s" % (PKG, version), "INFO")
            tgz = _fetch(tarball, timeout=300)
            # gzip 解包
            import gzip
            tar = gzip.decompress(tgz)
            bundle = _extract_tar_entry(tar, BUNDLE_ENTRY)
            if not bundle:
                raise RuntimeError("包内未找到 %s" % BUNDLE_ENTRY)
            elf = _find_elf(bundle, arch)
            tmp = BIN_PATH + ".tmp"
            with open(tmp, "wb") as f:
                f.write(elf)
            os.chmod(tmp, os.stat(tmp).st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH)
            os.replace(tmp, BIN_PATH)
            import hashlib
            sha = hashlib.sha256(elf).hexdigest()
            json.dump({"version": version, "arch": arch, "sha256": sha,
                       "installedAt": time.strftime("%Y-%m-%d %H:%M:%S"),
                       "registry": registry},
                      open(MANIFEST, "w", encoding="utf-8"), ensure_ascii=False, indent=1)
            log("umid: installed %s (arch=%s sha256=%s)" % (version, arch, sha[:12]), "INFO")
            return True, "已安装 %s（%s）" % (version, arch)
        except Exception as e:  # 尝试下一个 registry
            last_err = "%s: %s" % (registry, e)
            log("umid: registry %s failed: %s" % (registry, e), "WARN")
    return False, "下载失败（%s）" % last_err


# ---------------------------------------------------------------------------
# 调用
# ---------------------------------------------------------------------------

def run_runtime_info(env_code, account=""):
    """以指定 env 直接运行 runtime-info，返回 (完整输出 JSON, error)。

    env_code: 0=国内版 ONLINE，3=国际版（客户端默认），4=国际版（CLI 默认）。
    返回的 JSON 除 machineToken/Code/Type 外通常还有 vmInfo 等字段。
    """
    info = installed()
    if not info:
        return None, "未安装 runtime-info（请先在设置里下载）"
    try:
        stdin_data = json.dumps({"account": str(account or "")}) + "\n"
        p = subprocess.run([BIN_PATH, str(env_code), "--account-stdin"],
                           input=stdin_data, capture_output=True, text=True, timeout=30)
        if p.returncode != 0 or not p.stdout.strip():
            return None, "runtime-info 执行失败（退出码 %s）%s" % (p.returncode, (p.stderr or "")[:120])
        line = p.stdout.strip().split("\n")[0]
        data = json.loads(line)
        if not isinstance(data, dict):
            return None, "runtime-info 输出不是 JSON 对象"
        return data, None
    except subprocess.TimeoutExpired:
        return None, "runtime-info 执行超时"
    except Exception as e:
        return None, "runtime-info 调用异常：%s" % e


def run_runtime_info_raw(args, stdin_data=None):
    """按给定 argv / stdin 调用 runtime-info，返回 (JSON, error)。

    用途：对照不同的 env 取值与「是否带 --account-stdin」（客户端 Linux 分支
    只用 `runtime-info <env>`，完全不喂 stdin）。
    """
    info = installed()
    if not info:
        return None, "未安装 runtime-info（请先在设置里下载）"
    try:
        argv = [BIN_PATH] + [str(a) for a in args]
        if stdin_data is None:
            p = subprocess.run(argv, stdin=subprocess.DEVNULL,
                               capture_output=True, text=True, timeout=30)
        else:
            p = subprocess.run(argv, input=stdin_data,
                               capture_output=True, text=True, timeout=30)
        if p.returncode != 0 or not p.stdout.strip():
            return None, "退出码 %s %s" % (p.returncode, (p.stderr or "")[:160])
        line = p.stdout.strip().split("\n")[0]
        data = json.loads(line)
        if not isinstance(data, dict):
            return None, "输出不是 JSON 对象"
        return data, None
    except subprocess.TimeoutExpired:
        return None, "超时"
    except Exception as e:
        return None, "异常：%s" % e


def risk_identity(provider="qoder", account="", force_install=False):
    """生成设备风控身份。返回 ({"machineToken","machineCode","machineType"}, error).

    provider: "qoder"（国际版 env=4）/ "qoder-cn"（国内版 env=0）
    """
    info = installed()
    if not info or force_install:
        ok, msg = install(force=force_install)
        if not ok:
            return None, msg
        info = installed()
    data, err = run_runtime_info(RISK_ENV.get(provider, 0), account)
    if err:
        return None, err
    out = {}
    for k in ("machineToken", "machineCode", "machineType"):
        v = data.get(k)
        if not isinstance(v, str) or not v.strip():
            return None, "runtime-info 输出缺少 %s" % k
        out[k] = v.strip()
    if data.get("vmInfo"):
        out["vmInfo"] = data["vmInfo"]
    return out, None


def status():
    """供 UI 展示的状态。"""
    info = installed()
    return {
        "installed": bool(info),
        "path": BIN_PATH,
        "version": (info or {}).get("version"),
        "arch": (info or {}).get("arch"),
        "installedAt": (info or {}).get("installedAt"),
        "supported": _arch() in ELF_MACHINE,
        "cpu_arch": _arch(),
    }