#!/usr/bin/env python3
"""读取 APK 内 AndroidManifest.xml（二进制 AXML）的 versionName / versionCode。

为什么需要它（T-AND-013）：

文档 §11.2 要求「版本号单一来源 + 可核对」。本仓库原有的
`app/build.gradle.kts` 里那个 `verifyReleaseVersion` 任务只比**构建配置**
（`versionName` ↔ `project.yaml`）。但版本号真正呈现给系统与用户的地方是
**APK 的 AndroidManifest** —— 安装界面、应用信息页、`dumpsys package` 读的都是它。
配置全对、产物却是旧版（构建缓存、复用旧 APK、打包时拷错文件）属于同一类
「校验位置少一个，等于没校验」的失效；Windows 端 v0.2.1 就踩过（见该仓库
`tasks/017-pe-version.md`：配置文件全一致，exe 属性页显示 0.1.0）。

因此这里直接读**产物本身**。不依赖 Android SDK：AXML 的属性名就在字符串池里，
按属性名文本匹配即可，无需解析资源映射表，也不需要 aapt2。

用法：
    python3 build/check_apk_version.py <apk> [--expect-version X.Y.Z]
                                            [--expect-code N]
                                            [--print-code-for X.Y.Z]
退出码：0 通过 / 1 不一致 / 2 用法或文件问题。
"""

from __future__ import annotations

import struct
import sys
import zipfile

# AXML chunk 类型
RES_STRING_POOL = 0x0001
RES_XML_START_ELEMENT = 0x0102

# Res_value 的 dataType
TYPE_STRING = 0x03
TYPE_INT_DEC = 0x10
TYPE_INT_HEX = 0x11


class AxmlError(Exception):
    pass


def _parse_string_pool(data: bytes, offset: int) -> list[str]:
    """解析 ResStringPool：返回字符串列表。支持 UTF-8 与 UTF-16 两种编码。"""
    chunk_type, header_size, _chunk_size = struct.unpack_from("<HHI", data, offset)
    if chunk_type != RES_STRING_POOL:
        raise AxmlError("offset %d 处不是字符串池（type=0x%04x）" % (offset, chunk_type))

    string_count, _style_count, flags, strings_start, _styles_start = struct.unpack_from(
        "<IIIII", data, offset + 8
    )
    utf8 = bool(flags & (1 << 8))
    base = offset + strings_start

    offsets = struct.unpack_from("<%dI" % string_count, data, offset + header_size)
    out: list[str] = []
    for rel in offsets:
        pos = base + rel
        if utf8:
            # u16 字符数（高位为 1 表示后面还有一字节的字节数）、再 u8 字节数
            first = struct.unpack_from("<H", data, pos)[0]
            pos += 2
            char_len = first
            if first & 0x80:
                char_len = ((first & 0x7F) << 8) | data[pos]
                pos += 1
            byte_len = data[pos]
            pos += 1
            del char_len  # 只用字节数读取，避免两处长度不一致时误截断
            out.append(data[pos : pos + byte_len].decode("utf-8", "replace"))
        else:
            first = struct.unpack_from("<H", data, pos)[0]
            pos += 2
            char_len = first
            if first & 0x8000:
                char_len = ((first & 0x7FFF) << 16) | struct.unpack_from("<H", data, pos)[0]
                pos += 2
            out.append(data[pos : pos + char_len * 2].decode("utf-16-le", "replace"))
    return out


def read_manifest_versions(apk_path: str) -> tuple[str | None, int | None]:
    """读 <manifest> 标签的 versionName / versionCode。"""
    try:
        with zipfile.ZipFile(apk_path) as zf:
            raw = zf.read("AndroidManifest.xml")
    except (OSError, KeyError, zipfile.BadZipFile) as exc:
        raise AxmlError("无法读取 %s 里的 AndroidManifest.xml：%s" % (apk_path, exc)) from exc

    if len(raw) < 8:
        raise AxmlError("AndroidManifest.xml 太短（%d 字节），不是二进制 AXML" % len(raw))

    # 顶层 ResChunk_header（type=0x0003 XML），其后依次是字符串池、资源映射、XML 节点
    offset = struct.unpack_from("<HHI", raw, 0)[2]  # 用 header size 定位第一个子 chunk
    node_offset, header_size, _file_size = struct.unpack_from("<HHI", raw, 0)
    del node_offset
    if offset <= 0 or offset >= len(raw):
        offset = header_size if header_size else 8

    strings = _parse_string_pool(raw, offset)
    pos = offset + struct.unpack_from("<HHI", raw, offset)[2]

    while pos + 8 <= len(raw):
        chunk_type, _chunk_header_size, chunk_size = struct.unpack_from("<HHI", raw, pos)
        if chunk_size == 0:
            break

        if chunk_type == RES_XML_START_ELEMENT:
            # ResXMLTree_node(16) + ResXMLTree_attrExt(20) + attributeCount * 20
            (ns_idx, name_idx) = struct.unpack_from("<II", raw, pos + 16)
            attr_start, attr_size, attr_count = struct.unpack_from("<HHH", raw, pos + 16 + 8)
            tag = strings[name_idx] if name_idx < len(strings) else ""
            del ns_idx
            if tag == "manifest":
                version_name: str | None = None
                version_code: int | None = None
                attr_base = pos + 16 + attr_start
                for i in range(attr_count):
                    a = attr_base + i * attr_size
                    a_name_idx, a_raw_idx = struct.unpack_from("<II", raw, a + 4)
                    v_size, _res0, data_type, data = struct.unpack_from("<HBBI", raw, a + 12)
                    del v_size, _res0
                    a_name = strings[a_name_idx] if a_name_idx < len(strings) else ""
                    if a_name == "versionName":
                        if data_type == TYPE_STRING and data < len(strings):
                            version_name = strings[data]
                        elif a_raw_idx < len(strings):
                            version_name = strings[a_raw_idx]
                    elif a_name == "versionCode":
                        if data_type in (TYPE_INT_DEC, TYPE_INT_HEX):
                            version_code = data
                return version_name, version_code

        pos += chunk_size

    raise AxmlError("AndroidManifest.xml 里没有找到 <manifest> 标签")


def expected_version_code(version: str) -> int:
    """文档 §11.1：versionCode = 10000 * 主 + 100 * 次 + 修订。"""
    major, minor, patch = (int(x) for x in version.split("."))
    return 10000 * major + 100 * minor + patch


def main(argv: list[str]) -> int:
    if len(argv) < 2:
        print(__doc__.strip(), file=sys.stderr)
        return 2

    apk = argv[1]
    expect_version = None
    expect_code = None
    print_code_for = None

    i = 2
    while i < len(argv):
        arg = argv[i]
        if arg == "--expect-version" and i + 1 < len(argv):
            expect_version = argv[i + 1]
            i += 2
        elif arg == "--expect-code" and i + 1 < len(argv):
            expect_code = int(argv[i + 1])
            i += 2
        elif arg == "--print-code-for" and i + 1 < len(argv):
            print_code_for = argv[i + 1]
            i += 2
        else:
            print("未知参数：%s" % arg, file=sys.stderr)
            return 2

    if print_code_for is not None:
        print(expected_version_code(print_code_for))
        return 0

    try:
        version_name, version_code = read_manifest_versions(apk)
    except AxmlError as exc:
        print("✗ %s" % exc, file=sys.stderr)
        return 2

    print("  APK 产物版本元数据（%s）：" % apk)
    print("    AndroidManifest  versionName = %s" % (version_name if version_name else "（未找到）"))
    print("    AndroidManifest  versionCode = %s" % (version_code if version_code is not None else "（未找到）"))

    if expect_version is None and expect_code is None:
        return 0

    failed = False
    if expect_version is not None:
        if version_name != expect_version:
            print(
                "  ✗ APK 的 versionName 是 '%s'，期望 '%s'（安装界面与应用信息页会显示错版本）"
                % (version_name, expect_version),
                file=sys.stderr,
            )
            failed = True
        derived = expected_version_code(expect_version)
        if version_code is not None and version_code != derived:
            print(
                "  ✗ APK 的 versionCode 是 %s，期望 %s（§11.1：10000*主 + 100*次 + 修订）"
                % (version_code, derived),
                file=sys.stderr,
            )
            failed = True
    if expect_code is not None and version_code != expect_code:
        print("  ✗ APK 的 versionCode 是 %s，期望 %s" % (version_code, expect_code), file=sys.stderr)
        failed = True

    if failed:
        return 1
    print("  ✓ APK 版本元数据一致：%s（versionCode %s）" % (version_name, version_code))
    return 0


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