#!/usr/bin/env python3
"""Rebuild an Android expansion OBB from the matching Apktool-decoded tree."""
from __future__ import annotations

import argparse
import os
import sys
import tempfile
import zipfile
from pathlib import Path


def read_no_compress(config: Path) -> tuple[str, list[str]]:
    lines = config.read_text(encoding="utf-8").splitlines()
    name = next((line.split(":", 1)[1].strip().strip("'\"")
                 for line in lines if line.startswith("apkFileName:")), None)
    if not name:
        raise ValueError(f"{config} has no apkFileName")
    patterns: list[str] = []
    in_list = False
    for line in lines:
        if line == "doNotCompress:":
            in_list = True
            continue
        if in_list and line.startswith("- "):
            patterns.append(line[2:].strip().strip("'\"").replace("\\", "/"))
        elif in_list and line and not line.startswith(" "):
            break
    if not patterns:
        raise ValueError(f"{config} has no doNotCompress list")
    return name, patterns


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("decoded_obb", help="decoded OBB root containing apktool.yml and assets/")
    parser.add_argument("output", help="path for the rebuilt .obb")
    parser.add_argument("--check-only", action="store_true", help="validate inputs and print a packing summary")
    args = parser.parse_args()

    root = Path(args.decoded_obb).resolve()
    config = root / "apktool.yml"
    assets = root / "assets"
    if not config.is_file() or not assets.is_dir():
        print("decoded OBB must contain apktool.yml and assets/", file=sys.stderr)
        return 2
    output = Path(args.output).resolve()
    if output.exists() and not args.check_only:
        print(f"refusing to overwrite existing output: {output}", file=sys.stderr)
        return 2

    try:
        obb_name, patterns = read_no_compress(config)
        files = sorted((path for path in assets.rglob("*") if path.is_file()),
                       key=lambda path: path.relative_to(root).as_posix())
        names = {path.relative_to(root).as_posix(): path for path in files}
        exact = {pattern for pattern in patterns if pattern in names}
        extension_rules = {pattern.lower().lstrip(".") for pattern in patterns if "/" not in pattern and pattern not in names}
        unknown = [pattern for pattern in patterns if pattern not in names and pattern not in extension_rules]
        if unknown:
            raise ValueError(f"unmatched doNotCompress entries: {unknown[:8]}")
        stored = {name for name in names if name in exact or
                  Path(name).suffix.lower().lstrip(".") in extension_rules}
        matched_extensions = {Path(name).suffix.lower().lstrip(".") for name in names
                              if Path(name).suffix.lower().lstrip(".") in extension_rules}
        missing_extensions = sorted(extension_rules - matched_extensions)
        if missing_extensions:
            raise ValueError(f"doNotCompress extensions have no matching file: {missing_extensions}")
        uncompressed_bytes = sum(path.stat().st_size for name, path in names.items() if name in stored)
        total_bytes = sum(path.stat().st_size for path in files)
        if output.name != obb_name:
            raise ValueError(f"output filename must be {obb_name!r}, not {output.name!r}")
    except (OSError, ValueError) as exc:
        print(f"cannot package decoded OBB: {exc}", file=sys.stderr)
        return 2

    print(f"source: {root}")
    print(f"output: {output}")
    print(f"entries: {len(files):,}; uncompressed bytes: {total_bytes:,}")
    print(f"stored entries: {len(stored):,}; stored bytes: {uncompressed_bytes:,}; deflated entries: {len(files) - len(stored):,}")
    if args.check_only:
        return 0

    output.parent.mkdir(parents=True, exist_ok=True)
    fd, temporary_name = tempfile.mkstemp(prefix=output.name + ".", suffix=".tmp", dir=output.parent)
    os.close(fd)
    temporary = Path(temporary_name)
    try:
        with zipfile.ZipFile(temporary, "w", allowZip64=True, compression=zipfile.ZIP_DEFLATED,
                             compresslevel=6) as archive:
            for path in files:
                name = path.relative_to(root).as_posix()
                archive.write(path, arcname=name,
                              compress_type=zipfile.ZIP_STORED if name in stored else zipfile.ZIP_DEFLATED,
                              compresslevel=None if name in stored else 6)
        # Validate every CRC and path before publishing the archive at its final name.
        with zipfile.ZipFile(temporary, "r") as archive:
            bad_entry = archive.testzip()
            if bad_entry is not None:
                raise ValueError(f"archive CRC check failed at {bad_entry}")
            if set(archive.namelist()) != set(names):
                raise ValueError("archive entry list differs from the decoded assets tree")
        os.replace(temporary, output)
    except Exception:
        temporary.unlink(missing_ok=True)
        raise

    print(f"created {output} ({output.stat().st_size:,} bytes); CRC and entry-list checks passed")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())