"""从小X分身/虎牙 App 的 WebView Cookie 库提取虎牙 CK。

前提：手机已 root（SukiSU/KernelSU/Magisk），且 `adb shell su -c id` 返回 uid=0。
脚本只读设备数据，不修改 App 数据；默认只输出脱敏摘要，避免 CK 直接落到终端。

典型用法：
    .venv/bin/python scripts/huya_clone_ck.py
    .venv/bin/python scripts/huya_clone_ck.py --include-real --out /tmp/huya_cks.txt
    .venv/bin/python scripts/huya_clone_ck.py --verify-sid 2203

输出文件格式为 `账号----ck`，可直接贴入现有虎牙协作台账的「贴 / 换 CK」。
"""

from __future__ import annotations

import argparse
import json
import os
import shlex
import sqlite3
import subprocess
import sys
import tempfile
from collections.abc import Mapping
from dataclasses import dataclass, field
from datetime import UTC, datetime
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from core.huya.cookie_utils import cookie_pairs, cookie_value, normalize_cookie_pairs
from core.huya.web_cookie_fields import fill_web_cookie_fields

HUYA_PACKAGE = "com.duowan.kiwi"
DKPLAT_PACKAGE = "com.bly.dkplat"
COOKIE_REL = "app_webview/Default/Cookies"
CRED0_REL = "files/huyaudb/cred0"
# 虎牙 App 本地登录态文件 cred0 的 AES-128-ECB 密钥（UDB 主密钥前 16 字节）
UDB_MASTER_KEY = b"HuyaUdb1928374650qwertyuiop"[:16]
# WSS 登录态字段：App 登录 + cred0 就能给全
LOGIN_REQUIRED = ("udb_cred", "udb_biztoken", "udb_uid", "yyuid", "sdid")
# 匿名设备态字段：由匿名设备态协议签发（--anon-device / 后端自动补）
DEVICE_REQUIRED = ("udb_deviceid", "udb_guiddata", "udb_anobiztoken")
WSS_REQUIRED = LOGIN_REQUIRED + DEVICE_REQUIRED


class CliError(RuntimeError):
    """可预期的命令行错误。"""


@dataclass(frozen=True)
class CloneSource:
    label: str
    kind: str  # clone / real
    package: str
    data_dir: str
    cookies_path: str = ""


@dataclass
class Extracted:
    source: CloneSource
    uid: int
    passport: str
    cookie: str
    raw_cookie: str
    has_cred: bool
    missing_wss: tuple[str, ...]
    notes: list[str] = field(default_factory=list)
    verify: dict[str, object] | None = None

    @property
    def account(self) -> str:
        return self.passport or f"uid_{self.uid}"

    @property
    def wss_ready(self) -> bool:
        """登录态字段齐全即可走 WSS（匿名设备字段可另补）。"""
        present = {key for key, value in cookie_pairs(self.cookie) if value}
        return all(key in present for key in LOGIN_REQUIRED)


class AdbDevice:
    def __init__(self, adb: str, serial: str | None) -> None:
        self.adb = adb
        self.serial = serial

    def _invoke(
        self, mode: str, command: str, *, root: bool
    ) -> subprocess.CompletedProcess[bytes]:
        args = [self.adb]
        if self.serial:
            args.extend(["-s", self.serial])
        args.append(mode)
        args.append(f"su -c {shlex.quote(command)}" if root else command)
        try:
            return subprocess.run(args, capture_output=True, check=False)
        except FileNotFoundError as exc:
            raise CliError(f"找不到 adb 可执行文件: {self.adb}") from exc

    def shell_text(self, command: str, *, root: bool = False) -> str:
        proc = self._invoke("shell", command, root=root)
        if proc.returncode != 0:
            stderr = proc.stderr.decode("utf-8", "replace").strip()
            raise CliError(f"adb shell 失败（{command}）：{stderr or proc.returncode}")
        return proc.stdout.decode("utf-8", "replace")

    def exec_out(self, command: str, *, root: bool = False) -> bytes:
        proc = self._invoke("exec-out", command, root=root)
        if proc.returncode != 0 or not proc.stdout:
            stderr = proc.stderr.decode("utf-8", "replace").strip()
            raise CliError(f"adb exec-out 失败（{command}）：{stderr or '无输出'}")
        return proc.stdout

    def ensure_root(self) -> None:
        proc = self._invoke("shell", "id", root=True)
        output = (proc.stdout + proc.stderr).decode("utf-8", "replace")
        if proc.returncode != 0 or "uid=0" not in output:
            raise CliError(
                "adb shell 没有 root。请先在 SukiSU/KernelSU 里给 Shell 授权，"
                "确认 `adb shell su -c id` 返回 uid=0。"
            )


def pick_serial(adb: str, requested: str | None) -> str | None:
    if requested:
        return requested
    try:
        proc = subprocess.run([adb, "devices"], capture_output=True, check=False)
    except FileNotFoundError as exc:
        raise CliError(f"找不到 adb 可执行文件: {adb}") from exc
    serials = [
        line.split()[0]
        for line in proc.stdout.decode("utf-8", "replace").splitlines()
        if "\tdevice" in line and line.split()
    ]
    if not serials:
        raise CliError("没有已连接的 adb 设备。")
    if len(serials) > 1:
        raise CliError(
            "连接了多个 adb 设备，请用 --serial 指定目标：" + ", ".join(serials)
        )
    return serials[0]


def discover_sources(device: AdbDevice, *, include_real: bool) -> list[CloneSource]:
    sources: list[CloneSource] = []

    packages = device.shell_text("pm list packages -3")
    for line in packages.splitlines():
        package = line.strip().removeprefix("package:")
        if not package.startswith("dkplugin."):
            continue
        sources.append(
            CloneSource(
                label=f"clone:{package.removeprefix('dkplugin.')}",
                kind="clone",
                package=package,
                data_dir=f"/data/user/0/{package}/virtual/data/user/0/{HUYA_PACKAGE}",
            )
        )

    # 小X分身自己的「主分身」数据在 com.bly.dkplat/chaos/data/user/<id>/ 下。
    glob_cmd = (
        f"ls -d /data/user/0/{DKPLAT_PACKAGE}/chaos/data/user/*/{HUYA_PACKAGE} "
        "2>/dev/null || true"
    )
    for raw_path in device.shell_text(glob_cmd, root=True).splitlines():
        path = raw_path.strip()
        if not path:
            continue
        marker = f"/{DKPLAT_PACKAGE}/chaos/data/user/"
        instance = path.split(marker, 1)[1].split("/", 1)[0] if marker in path else "?"
        sources.append(
            CloneSource(
                label=f"clone:dkplat/{instance}",
                kind="clone",
                package=DKPLAT_PACKAGE,
                data_dir=path,
            )
        )

    if include_real:
        sources.append(
            CloneSource(
                label="real:huya",
                kind="real",
                package=HUYA_PACKAGE,
                data_dir=f"/data/user/0/{HUYA_PACKAGE}",
            )
        )

    found: list[CloneSource] = []
    for source in sources:
        cookies_path = f"{source.data_dir}/{COOKIE_REL}"
        probe = device.shell_text(
            f"test -f {cookies_path} && echo yes || echo no", root=True
        ).strip()
        if probe.endswith("yes"):
            found.append(
                CloneSource(
                    label=source.label,
                    kind=source.kind,
                    package=source.package,
                    data_dir=source.data_dir,
                    cookies_path=cookies_path,
                )
            )
    return sorted(found, key=lambda item: item.label)


def read_cookies(blob: bytes) -> tuple[str, list[str]]:
    """把 WebView Cookies SQLite 转成 Cookie 字符串。"""
    with tempfile.NamedTemporaryFile(suffix=".db") as handle:
        handle.write(blob)
        handle.flush()
        connection = sqlite3.connect(f"file:{handle.name}?mode=ro", uri=True)
        try:
            integrity = connection.execute("PRAGMA integrity_check").fetchone()
            if not integrity or integrity[0] != "ok":
                raise CliError(f"Cookie 数据库完整性检查失败: {integrity}")
            rows = list(
                connection.execute(
                    "SELECT host_key, name, value, encrypted_value "
                    "FROM cookies ORDER BY rowid"
                )
            )
        finally:
            connection.close()

    pairs: list[tuple[str, str]] = []
    encrypted: list[str] = []
    for host_key, name, value, encrypted_value in rows:
        host = str(host_key or "")
        key = str(name or "")
        if not key or not (host == ".huya.com" or host.endswith("huya.com")):
            continue
        if value:
            pairs.append((key, str(value)))
        elif encrypted_value:
            encrypted.append(key)
    if not pairs:
        raise CliError("Cookie 库里没有非空的 huya.com Cookie。")
    return normalize_cookie_pairs(pairs), encrypted


def decrypt_cred0(blob: bytes) -> dict[str, object] | None:
    """解密 ``files/huyaudb/cred0``，取回 App 登录态里的 cred/uid/passport。

    cred0 是 AES-128-ECB(零填充) 加密的登录历史 JSON；``apploginData.cred``
    就是 ``udb_cred``（152 字符 base64url），WebView Cookie 里永远没有它。
    解密失败只返回 None，不影响 Cookie 提取。
    """
    if not blob or len(blob) % 16:
        return None
    try:
        from Crypto.Cipher import AES

        plain = AES.new(UDB_MASTER_KEY, AES.MODE_ECB).decrypt(blob).rstrip(b"\x00")
        payload = json.loads(plain)
        history = payload.get("loginHistory") or []
        if not history:
            return None
        login = history[0].get("apploginData") or {}
        cred = str(login.get("cred") or "")
        if not cred:
            return None
        cookie = login.get("cookie") or {}
        return {
            "cred": cred,
            "uid": int(login.get("uid") or cookie.get("uid") or 0),
            "passport": str(login.get("passport") or cookie.get("passport") or ""),
            "biztoken": str(cookie.get("biztoken") or ""),
            "status": str(cookie.get("status") or ""),
            "version": str(cookie.get("version") or ""),
        }
    except Exception:  # noqa: BLE001 - 解密失败按「没有 cred0」处理
        return None


def apply_credential(raw_cookie: str, login: Mapping[str, object]) -> str:
    """把 cred0 里的登录凭据写进 Cookie（覆盖 WebView 里的旧值）。"""
    overrides = {
        "udb_cred": login.get("cred"),
        "udb_uid": login.get("uid"),
        "yyuid": login.get("uid"),
        "udb_passport": login.get("passport"),
        "username": login.get("passport"),
        "udb_biztoken": login.get("biztoken"),
        "udb_status": login.get("status"),
        "udb_version": login.get("version"),
    }
    pairs = [
        (key, value) for key, value in cookie_pairs(raw_cookie) if key not in overrides
    ]
    for key, value in overrides.items():
        text = "" if value in (None, 0) else str(value)
        if text:
            pairs.append((key, text))
    return "; ".join(f"{key}={value}" for key, value in pairs)


def build_web_cookie(raw_cookie: str, uid: int) -> str:
    """补齐网页设备态，并把渠道 appId 修正为网页的 5002。"""
    filled = fill_web_cookie_fields(raw_cookie, guid="", uid=uid, qimei_feature=None)
    pairs = dict(cookie_pairs(filled))
    pairs["udb_appid"] = "5002"
    return "; ".join(f"{key}={value}" for key, value in pairs.items())


def verify_cookie(cookie: str, uid: int, sid: int, timeout: float) -> dict[str, object]:
    from core.huya.http_client import HuyaHttpClient

    try:
        response = HuyaHttpClient(logger=lambda _message: None).query_user_score(
            uid, cookie, sid=sid, timeout=timeout
        )
    except Exception as exc:  # noqa: BLE001 - 验证失败要写进报告而不是中断提取
        return {"ok": False, "error": f"{type(exc).__name__}: {exc}"}
    return {
        "ok": True,
        "status": getattr(response, "status", None),
        "msg": getattr(response, "msg", None),
    }


def extract_one(
    device: AdbDevice,
    source: CloneSource,
    *,
    fill_web: bool,
    verify_sid: int | None,
    timeout: float,
    use_cred0: bool = True,
    anon_device: bool = False,
) -> Extracted:
    blob = device.exec_out(f"cat {source.cookies_path}", root=True)
    raw_cookie, encrypted = read_cookies(blob)
    uid = int(cookie_value(raw_cookie, "udb_uid") or 0)
    if not uid:
        uid = int(cookie_value(raw_cookie, "yyuid") or 0)

    notes: list[str] = []
    if encrypted:
        notes.append(
            f"有 {len(encrypted)} 个 Cookie 使用 encrypted_value（当前脚本不解密）"
        )

    # WebView Cookie 只有登录态字段，udb_cred 只存在于 App 本地登录态 cred0。
    login: dict[str, object] | None = None
    if use_cred0:
        cred_path = f"{source.data_dir}/{CRED0_REL}"
        probe = device.shell_text(
            f"test -f {cred_path} && echo yes || echo no", root=True
        ).strip()
        if probe.endswith("yes"):
            login = decrypt_cred0(device.exec_out(f"cat {cred_path}", root=True))
            if login:
                notes.append("已从 cred0 解出 App 登录凭据（udb_cred）")
                if not uid and login.get("uid"):
                    uid = int(login["uid"])  # type: ignore[arg-type]
            else:
                notes.append("cred0 存在但解密失败，CK 不含 udb_cred")
        else:
            notes.append("没有 cred0（该分身可能只有匿名态）")

    if not uid:
        raise CliError(f"{source.label} 的 Cookie 里没有 udb_uid/yyuid。")

    base_cookie = apply_credential(raw_cookie, login) if login else raw_cookie

    try:
        cookie = build_web_cookie(base_cookie, uid) if fill_web else base_cookie
    except Exception as exc:  # noqa: BLE001 - 补齐失败时保留原始 Cookie
        notes.append(
            f"网页设备态补齐失败，退回原始 Cookie：{type(exc).__name__}: {exc}"
        )
        cookie = base_cookie

    if anon_device:
        try:
            from core.huya.anon_device import fetch_anonymous_device_state

            anon = fetch_anonymous_device_state()
            missing = {
                key: value
                for key, value in anon.items()
                if value and not cookie_value(cookie, key)
            }
            if missing:
                cookie = "; ".join(
                    [cookie] + [f"{key}={value}" for key, value in missing.items()]
                )
                notes.append("已用匿名设备态补齐: " + ", ".join(sorted(missing)))
        except Exception as exc:  # noqa: BLE001 - 匿名设备态失败不阻断提取
            notes.append(f"匿名设备态获取失败: {type(exc).__name__}: {exc}")

    missing_login = tuple(
        key for key in LOGIN_REQUIRED if not cookie_value(cookie, key)
    )
    missing_device = tuple(
        key for key in DEVICE_REQUIRED if not cookie_value(cookie, key)
    )
    missing_wss = missing_login + missing_device
    if missing_login:
        notes.append(
            "缺少 WSS 登录字段: "
            + ", ".join(missing_login)
            + "；活动/商城只能走 HTTP 通道"
        )
    elif missing_device:
        notes.append(
            "缺少匿名设备字段: "
            + ", ".join(missing_device)
            + "（加 --anon-device 可由协议补齐）"
        )

    verify = (
        verify_cookie(cookie, uid, verify_sid, timeout)
        if verify_sid is not None
        else None
    )
    return Extracted(
        source=source,
        uid=uid,
        passport=cookie_value(raw_cookie, "udb_passport")
        or cookie_value(raw_cookie, "username"),
        cookie=cookie,
        raw_cookie=raw_cookie,
        has_cred=bool(cookie_value(cookie, "udb_cred")),
        missing_wss=missing_wss,
        notes=notes,
        verify=verify,
    )


def public_record(item: Extracted, *, reveal: bool) -> dict[str, object]:
    record: dict[str, object] = {
        "label": item.source.label,
        "kind": item.source.kind,
        "package": item.source.package,
        "uid": item.uid,
        "passport": item.passport,
        "account": item.account,
        "cookie_fields": len(cookie_pairs(item.cookie)),
        "cookie_bytes": len(item.cookie.encode("utf-8")),
        "has_cred": item.has_cred,
        "wss_ready": item.wss_ready,
        "missing_device": [k for k in DEVICE_REQUIRED if k in item.missing_wss],
        "missing_wss": list(item.missing_wss),
        "notes": item.notes,
    }
    if reveal:
        record["cookie"] = item.cookie
    if item.verify is not None:
        record["verify"] = item.verify
    return record


def print_table(items: list[Extracted]) -> None:
    if not items:
        print("没有发现可提取的虎牙 Cookie。")
        return
    headers = ("实例", "UID", "账号", "字段", "WSS", "验证", "说明")
    rows = []
    for item in items:
        verify = "-"
        if item.verify is not None:
            if item.verify.get("ok"):
                verify = f"{item.verify.get('status')} {item.verify.get('msg') or ''}".strip()
            else:
                verify = str(item.verify.get("error") or "失败")
        rows.append(
            (
                item.source.label,
                str(item.uid),
                item.account,
                str(len(cookie_pairs(item.cookie))),
                "OK" if item.wss_ready else "HTTP",
                verify,
                "；".join(item.notes) or "-",
            )
        )
    widths = [
        max(len(headers[index]), *(len(row[index]) for row in rows))
        for index in range(len(headers))
    ]
    print(
        "  ".join(headers[index].ljust(widths[index]) for index in range(len(headers)))
    )
    for row in rows:
        print(
            "  ".join(row[index].ljust(widths[index]) for index in range(len(headers)))
        )


def write_cookie_file(path: Path, items: list[Extracted]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    lines = [f"{item.account}----{item.cookie}" for item in items]
    path.write_text("\n".join(lines) + "\n", encoding="utf-8")
    os.chmod(path, 0o600)


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="从已 root 的安卓设备提取小X分身/虎牙 App 的 CK（只读）"
    )
    parser.add_argument("--adb", default="adb", help="adb 可执行文件（默认: adb）")
    parser.add_argument("--serial", help="目标设备序列号（多设备时必填）")
    parser.add_argument(
        "--include-real",
        action="store_true",
        help="同时提取真身 com.duowan.kiwi（默认只提分身）",
    )
    parser.add_argument(
        "--raw",
        action="store_true",
        help="不补齐网页设备字段（仍会注入 cred0 解出的 udb_cred）",
    )
    parser.add_argument(
        "--no-cred0",
        action="store_true",
        help="不解密 files/huyaudb/cred0（默认会解密并补上 udb_cred）",
    )
    parser.add_argument(
        "--anon-device",
        action="store_true",
        help="额外走匿名设备态协议补齐 udb_deviceid/guiddata/anobiztoken（需要外网）",
    )
    parser.add_argument(
        "--out",
        type=Path,
        help="把完整 CK 写成 `账号----ck` 文件（权限 0600）",
    )
    parser.add_argument(
        "--print",
        action="store_true",
        help="在终端打印完整 `账号----ck`（默认只打印脱敏摘要）",
    )
    parser.add_argument("--json", action="store_true", help="输出 JSON 摘要")
    parser.add_argument(
        "--verify-sid",
        type=int,
        help="可选：用活动积分接口验证 CK（例如 2203；活动已结束也会返回 603）",
    )
    parser.add_argument("--timeout", type=float, default=15.0, help="验证超时秒数")
    return parser


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    serial = pick_serial(args.adb, args.serial)
    device = AdbDevice(args.adb, serial)
    device.ensure_root()

    sources = discover_sources(device, include_real=args.include_real)
    if not sources:
        raise CliError(
            "没有发现虎牙分身。请确认小X分身已安装虎牙、分身至少启动过一次。"
        )

    items: list[Extracted] = []
    seen: dict[tuple[int, str], str] = {}
    for source in sources:
        try:
            item = extract_one(
                device,
                source,
                fill_web=not args.raw,
                verify_sid=args.verify_sid,
                timeout=args.timeout,
                use_cred0=not args.no_cred0,
                anon_device=args.anon_device,
            )
        except CliError as exc:
            print(f"[跳过] {source.label}: {exc}", file=sys.stderr)
            continue
        key = (item.uid, item.passport)
        if key in seen:
            print(
                f"[去重] {source.label} 与 {seen[key]} 是同一账号 {item.account}",
                file=sys.stderr,
            )
            continue
        seen[key] = source.label
        items.append(item)

    if args.out:
        write_cookie_file(args.out, items)
    if args.print:
        for item in items:
            print(f"{item.account}----{item.cookie}")

    if args.json:
        print(
            json.dumps(
                [public_record(item, reveal=args.print) for item in items],
                ensure_ascii=False,
                indent=2,
            )
        )
    elif not args.print:
        print_table(items)

    if args.out:
        print(
            f"已写入 {len(items)} 个 CK: {args.out}（{datetime.now(UTC).isoformat()}）",
            file=sys.stderr,
        )
    if not args.print and not args.out and not args.json:
        print("提示：加 --out 落盘，或加 --print 在终端输出完整 CK。")
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except CliError as exc:
        print(f"错误：{exc}", file=sys.stderr)
        raise SystemExit(2) from exc