#!/usr/bin/env python3
"""Build deterministic, map-scoped PK4 bundles from native OpenQ4 I/O traces."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import pathlib
import re
import shutil
import struct
import sys
import tempfile
import zipfile
from collections import deque
from dataclasses import dataclass
from typing import Iterable


RETAIL_SCHEMA = 1
RETAIL_POLICY = "operator-authorized-retail-v1"
PROFILE_SCHEMA = 1
PROFILE_POLICY = "openq4-fs-debug-working-set-v1"
BUNDLE_SCHEMA = 1
BUNDLE_POLICY = "operator-authorized-retail-bundles-v1"
PERFORMANCE_PRESET = "balanced"
MAX_MACHINE_SPEC = 2
FIXED_TIME = (1980, 1, 1, 0, 0, 0)
CHUNK_SIZE = 1024 * 1024
MAP_PATTERN = re.compile(r"^mp/[a-z0-9_]+$")
SHA256_PATTERN = re.compile(r"^[a-f0-9]{64}$")
FS_DEBUG_PATTERN = re.compile(
    r"idFileSystem::OpenFileRead:\s+(.+?)\s+\(found in '(.*?)'\)"
)
REFERENCE_PATTERN = re.compile(r"[A-Za-z0-9_./\\-]+")
EXPECTED_RETAIL = tuple(
    [f"pak{index:03}.pk4" for index in range(1, 23)] + ["zpak_english.pk4"]
)
REQUIRED_COVERAGE = frozenset(
    {
        "fresh-process-per-map",
        "balanced-machine-spec-2",
        "all-player-models",
        "all-mp-weapons-items-projectiles",
        "dm-tourney-team-ctf-gui-audio",
    }
)


@dataclass(frozen=True)
class Winner:
    key: str
    name: str
    archive_name: str
    archive_path: pathlib.Path
    member_name: str
    size: int
    crc: int


@dataclass(frozen=True)
class AssetToken:
    value: str
    quoted: bool
    start: int
    end: int


@dataclass(frozen=True)
class GuideTemplate:
    key: str
    parameters: tuple[str, ...]
    body: str


@dataclass(frozen=True)
class Declaration:
    key: str
    name: str
    body: str
    support_keys: tuple[str, ...] = ()


class BuildError(ValueError):
    pass


def sha256_file(file: pathlib.Path) -> str:
    digest = hashlib.sha256()
    with file.open("rb") as stream:
        while chunk := stream.read(CHUNK_SIZE):
            digest.update(chunk)
    return digest.hexdigest()


def normalize_member(name: str) -> tuple[str, str]:
    clean = name.replace("\\", "/")
    parts = clean.split("/")
    if not clean or clean.startswith("/") or "\0" in clean or any(
        part in {"", ".", ".."} for part in parts
    ) or re.match(r"^[a-zA-Z]:", clean):
        raise BuildError(f"unsafe PK4 member path: {name!r}")
    return clean.casefold(), clean


def read_json(file: pathlib.Path) -> dict:
    try:
        value = json.loads(file.read_text(encoding="utf-8"))
    except (OSError, UnicodeError, json.JSONDecodeError) as error:
        raise BuildError(f"cannot read JSON {file}: {error}") from error
    if not isinstance(value, dict):
        raise BuildError(f"JSON root must be an object: {file}")
    return value


def canonical_retail_fingerprint(files: Iterable[dict]) -> str:
    digest = hashlib.sha256()
    for entry in sorted(files, key=lambda item: item["name"]):
        digest.update(
            f"{entry['name']}\0{entry['size']}\0{entry['sha256']}\n".encode("utf-8")
        )
    return digest.hexdigest()


def load_retail(retail_dir: pathlib.Path, manifest_file: pathlib.Path) -> tuple[dict, list[dict]]:
    manifest = read_json(manifest_file)
    if manifest.get("schema") != RETAIL_SCHEMA or manifest.get("policy") != RETAIL_POLICY:
        raise BuildError("retail manifest has the wrong schema or policy")
    raw_files = manifest.get("files")
    if not isinstance(raw_files, list):
        raise BuildError("retail manifest files must be an array")
    by_name = {str(entry.get("name", "")).lower(): entry for entry in raw_files if isinstance(entry, dict)}
    if set(by_name) != set(EXPECTED_RETAIL) or len(raw_files) != len(EXPECTED_RETAIL):
        raise BuildError("retail manifest must contain the exact English 23-file catalog")
    files: list[dict] = []
    for name in EXPECTED_RETAIL:
        raw = by_name[name]
        size = raw.get("size")
        digest = str(raw.get("sha256", "")).lower()
        source = retail_dir / name
        if not isinstance(size, int) or size <= 0 or not SHA256_PATTERN.fullmatch(digest):
            raise BuildError(f"invalid retail record: {name}")
        if not source.is_file() or source.is_symlink() or source.stat().st_size != size:
            raise BuildError(f"retail source is missing or has the wrong size: {source}")
        actual = sha256_file(source)
        if actual != digest:
            raise BuildError(f"retail source SHA-256 mismatch: {name}")
        files.append({"name": name, "size": size, "sha256": digest, "path": source})
    fingerprint = canonical_retail_fingerprint(files)
    if manifest.get("fingerprint") != fingerprint or manifest.get("revision") != f"q4retail-{fingerprint[:20]}":
        raise BuildError("retail revision or fingerprint is not canonical")
    return manifest, files


def build_override_graph(files: list[dict]) -> dict[str, Winner]:
    winners: dict[str, Winner] = {}
    for source in files:
        seen: set[str] = set()
        try:
            with zipfile.ZipFile(source["path"]) as archive:
                for member in archive.infolist():
                    if member.is_dir():
                        continue
                    key, name = normalize_member(member.filename)
                    if key in seen:
                        raise BuildError(f"case-insensitive duplicate inside {source['name']}: {name}")
                    seen.add(key)
                    winners[key] = Winner(
                        key=key,
                        name=name,
                        archive_name=source["name"],
                        archive_path=source["path"],
                        member_name=member.filename,
                        size=member.file_size,
                        crc=member.CRC,
                    )
        except zipfile.BadZipFile as error:
            raise BuildError(f"invalid retail PK4 {source['name']}: {error}") from error
    return winners


def lex_asset_text(contents: str, source: str) -> list[AssetToken]:
    tokens: list[AssetToken] = []
    index = 0
    length = len(contents)
    punctuation = "{}(),"
    while index < length:
        char = contents[index]
        if char.isspace():
            index += 1
            continue
        if contents.startswith("//", index):
            newline = contents.find("\n", index + 2)
            index = length if newline < 0 else newline + 1
            continue
        if contents.startswith("/*", index):
            end = contents.find("*/", index + 2)
            if end < 0:
                raise BuildError(f"unterminated block comment in {source}")
            index = end + 2
            continue
        if char == '"':
            start = index
            index += 1
            value: list[str] = []
            while index < length:
                char = contents[index]
                if char == '"':
                    index += 1
                    tokens.append(AssetToken("".join(value), True, start, index))
                    break
                if char == "\\" and index + 1 < length and contents[index + 1] in {'"', "\\"}:
                    value.append(contents[index + 1])
                    index += 2
                    continue
                value.append(char)
                index += 1
            else:
                raise BuildError(f"unterminated string in {source}")
            continue
        if char in punctuation:
            tokens.append(AssetToken(char, False, index, index + 1))
            index += 1
            continue
        start = index
        while index < length:
            char = contents[index]
            if (
                char.isspace()
                or char in punctuation + '"'
                or contents.startswith("//", index)
                or contents.startswith("/*", index)
            ):
                break
            index += 1
        if index == start:
            raise BuildError(f"cannot tokenize dependency source {source}")
        tokens.append(AssetToken(contents[start:index], False, start, index))
    return tokens


def matching_token(
    tokens: list[AssetToken],
    opening: int,
    left: str,
    right: str,
    source: str,
    *,
    allow_unclosed: bool = False,
) -> int:
    depth = 0
    for index in range(opening, len(tokens)):
        value = tokens[index].value
        if value == left:
            depth += 1
        elif value == right:
            depth -= 1
            if depth == 0:
                return index
    if allow_unclosed:
        return len(tokens)
    raise BuildError(f"unbalanced {left}{right} block in {source}")


def split_arguments(tokens: list[AssetToken]) -> tuple[str, ...]:
    arguments: list[list[str]] = [[]]
    depth = 0
    for token in tokens:
        if token.value == "(":
            depth += 1
        elif token.value == ")":
            depth -= 1
        if token.value == "," and depth == 0:
            arguments.append([])
        else:
            arguments[-1].append(token.value)
    return tuple(" ".join(argument).strip() for argument in arguments)


class RetailDependencyClosure:
    """Conservatively close map assets without precaching every startup declaration."""

    DECLARATION_PREFIXES = {
        ".af": frozenset({"articulatedfigure"}),
        ".def": frozenset({"entitydef", "export", "mapdef", "model"}),
        ".fx": frozenset({"effect"}),
        ".mtr": frozenset({"particle", "skin", "table"}),
        ".reverb": frozenset({"reverb"}),
        ".skin": frozenset({"skin"}),
        ".sndshd": frozenset({"sound"}),
    }
    DECLARATION_EXTENSIONS = frozenset(DECLARATION_PREFIXES)
    INITIAL_TRAVERSABLE = frozenset({".ase", ".cm", ".gui", ".lwo", ".map", ".md5mesh", ".proc"})
    REFERENCED_TRAVERSABLE = INITIAL_TRAVERSABLE | frozenset({".af", ".fx", ".skin"})
    IMAGE_EXTENSIONS = (".dds", ".tga")
    CUBE_SUFFIXES = (
        "_nx", "_ny", "_nz", "_px", "_py", "_pz",
        "_back", "_down", "_forward", "_left", "_right", "_up",
    )
    KNOWN_PATH_PREFIXES = (
        "effects/", "fonts/", "gfx/", "glprogs/", "guis/", "maps/",
        "models/", "progimg/", "sound/", "textures/", "video/",
    )

    def __init__(self, winners: dict[str, Winner]):
        self.winners = winners
        self._small_contents: dict[str, bytes] = {}
        self._archives: dict[pathlib.Path, zipfile.ZipFile] = {}
        self._reference_cache: dict[object, tuple[tuple[str, ...], tuple[str, ...]]] = {}
        self.guides: dict[str, GuideTemplate] = {}
        self.declarations: dict[str, list[Declaration]] = {}
        try:
            self._index_guides()
            self._index_declarations()
        except Exception:
            self.release()
            raise

    def release(self) -> None:
        for archive in self._archives.values():
            archive.close()
        self._archives.clear()

    def read_bytes(self, key: str) -> bytes:
        cached = self._small_contents.get(key)
        if cached is not None:
            return cached
        winner = self.winners[key]
        archive = self._archives.get(winner.archive_path)
        if archive is None:
            archive = zipfile.ZipFile(winner.archive_path)
            self._archives[winner.archive_path] = archive
        contents = archive.read(winner.member_name)
        if len(contents) <= CHUNK_SIZE:
            self._small_contents[key] = contents
        return contents

    def read_text(self, key: str) -> str:
        return self.read_bytes(key).decode("latin1", errors="replace")

    def _index_guides(self) -> None:
        for key in sorted(self.winners):
            if pathlib.PurePosixPath(key).suffix != ".guide":
                continue
            source = self.winners[key].name
            contents = self.read_text(key)
            tokens = lex_asset_text(contents, source)
            index = 0
            while index < len(tokens):
                if tokens[index].value.casefold() not in {"guide", "inlineguide"}:
                    index += 1
                    continue
                if index + 3 >= len(tokens) or tokens[index + 2].value != "(":
                    raise BuildError(f"invalid guide declaration in {source}")
                close_parameters = matching_token(tokens, index + 2, "(", ")", source)
                if close_parameters + 1 >= len(tokens) or tokens[close_parameters + 1].value != "{":
                    raise BuildError(f"guide has no body in {source}: {tokens[index + 1].value}")
                close_body = matching_token(tokens, close_parameters + 1, "{", "}", source)
                parameters = tuple(
                    token.value
                    for token in tokens[index + 3:close_parameters]
                    if token.value != ","
                )
                name = tokens[index + 1].value.casefold()
                template = GuideTemplate(
                    key=key,
                    parameters=parameters,
                    body=contents[tokens[close_parameters + 1].end:tokens[close_body].start],
                )
                previous = self.guides.setdefault(name, template)
                if previous != template:
                    raise BuildError(f"duplicate guide template: {tokens[index + 1].value}")
                index = close_body + 1

    def _add_declaration(self, declaration: Declaration) -> None:
        records = self.declarations.setdefault(declaration.name.casefold(), [])
        if declaration not in records:
            records.append(declaration)

    def _expanded_guide(
        self,
        key: str,
        material_name: str,
        guide_name: str,
        arguments: tuple[str, ...],
    ) -> Declaration:
        template = self.guides.get(guide_name.casefold())
        if template is None:
            return Declaration(key=key, name=material_name, body=" ".join(arguments))
        if len(arguments) != len(template.parameters):
            raise BuildError(
                f"guide argument mismatch for {material_name}: "
                f"{guide_name} expects {len(template.parameters)}, got {len(arguments)}"
            )
        body = template.body
        for parameter, argument in zip(template.parameters, arguments):
            body = re.sub(
                rf"(?<![A-Za-z0-9_]){re.escape(parameter)}(?=[^A-Za-z0-9]|$)",
                lambda _match, replacement=argument: replacement,
                body,
                flags=re.IGNORECASE,
            )
        return Declaration(
            key=key,
            name=material_name,
            body=body,
            support_keys=(template.key,),
        )

    def _index_declarations(self) -> None:
        for key in sorted(self.winners):
            extension = pathlib.PurePosixPath(key).suffix
            if extension not in self.DECLARATION_EXTENSIONS:
                continue
            source = self.winners[key].name
            contents = self.read_text(key)
            tokens = lex_asset_text(contents, source)
            prefixes = self.DECLARATION_PREFIXES[extension]
            index = 0
            while index < len(tokens):
                token = tokens[index]
                if extension == ".mtr" and token.value.casefold() == "guide":
                    if index + 3 >= len(tokens) or tokens[index + 3].value != "(":
                        raise BuildError(f"invalid material guide invocation in {source}")
                    close = matching_token(tokens, index + 3, "(", ")", source)
                    arguments = split_arguments(tokens[index + 4:close])
                    self._add_declaration(self._expanded_guide(
                        key,
                        tokens[index + 1].value,
                        tokens[index + 2].value,
                        arguments,
                    ))
                    index = close + 1
                    continue
                opening = -1
                name = ""
                if token.value.casefold() in prefixes:
                    if index + 2 < len(tokens) and tokens[index + 2].value == "{":
                        name = tokens[index + 1].value
                        opening = index + 2
                elif extension == ".mtr" and index + 1 < len(tokens) and tokens[index + 1].value == "{":
                    name = token.value
                    opening = index + 1
                if opening < 0:
                    index += 1
                    continue
                close = matching_token(
                    tokens,
                    opening,
                    "{",
                    "}",
                    source,
                    allow_unclosed=True,
                )
                body_end = len(contents) if close == len(tokens) else tokens[close].start
                self._add_declaration(Declaration(
                    key=key,
                    name=name,
                    body=contents[tokens[opening].end:body_end],
                ))
                index = close + 1

    def close(self, seeds: set[str]) -> set[str]:
        result = set(seeds)
        pending_files: deque[str] = deque()
        pending_declarations: deque[Declaration] = deque()
        traversed: set[str] = set()
        activated: set[Declaration] = set()

        def add_key(key: str, traverse: bool = False) -> None:
            if key not in self.winners:
                return
            result.add(key)
            if traverse and key not in traversed:
                pending_files.append(key)

        def add_image(raw: str) -> None:
            normalized = raw.replace("\\", "/").casefold()
            suffix = pathlib.PurePosixPath(normalized).suffix
            if suffix in self.IMAGE_EXTENSIONS:
                stem = normalized[:-len(suffix)]
                for extension in self.IMAGE_EXTENSIONS:
                    add_key(stem + extension)
            elif not suffix:
                for extension in self.IMAGE_EXTENSIONS:
                    add_key(normalized + extension)

        def add_cube(raw: str) -> None:
            normalized = raw.replace("\\", "/").casefold()
            for suffix in self.CUBE_SUFFIXES:
                add_image(normalized + suffix)

        def activate(raw: str) -> None:
            for declaration in self.declarations.get(raw.casefold(), ()):
                if declaration not in activated:
                    pending_declarations.append(declaration)

        def resolve_reference(raw: str) -> None:
            normalized = raw.replace("\\", "/").strip().casefold()
            if not normalized or normalized in {".", ".."}:
                return
            path_offsets = [
                normalized.find(prefix)
                for prefix in self.KNOWN_PATH_PREFIXES
                if normalized.find(prefix) >= 0
            ]
            if path_offsets and min(path_offsets) > 0:
                normalized = normalized[min(path_offsets):]
            activate(normalized)
            suffix = pathlib.PurePosixPath(normalized).suffix
            if normalized in self.winners:
                add_key(normalized, suffix in self.REFERENCED_TRAVERSABLE)
            if suffix in self.IMAGE_EXTENSIONS:
                add_image(normalized)
            elif suffix in {".ogg", ".wav"}:
                stem = normalized[:-len(suffix)]
                add_key(stem + ".ogg")
                add_key(stem + ".wav")
            elif not suffix and "/" in normalized:
                if normalized.startswith(("effects/", "fx/")):
                    add_key(normalized + ".fx", True)
                if normalized.startswith("guis/"):
                    add_key(normalized + ".gui", True)
                if normalized.startswith("sound/"):
                    add_key(normalized + ".ogg")
                    add_key(normalized + ".wav")
                if normalized.startswith(("gfx/", "models/", "textures/")):
                    add_image(normalized)
                if normalized.startswith("models/"):
                    for extension in (".ase", ".cm", ".lwo", ".md5mesh"):
                        add_key(normalized + extension, True)
                if normalized.startswith("fonts/"):
                    for size in (12, 24, 48):
                        add_key(f"{normalized}_{size}.fontdat")

        def scan(contents: str, identity: object) -> None:
            cached = self._reference_cache.get(identity)
            if cached is None:
                try:
                    comment_free = []
                    for token in lex_asset_text(contents, "<dependency body>"):
                        if token.value not in {"{", "}", "(", ")", ","}:
                            comment_free.append(token.value)
                    flattened = " ".join(comment_free)
                except BuildError:
                    # Some original binary/text hybrid assets contain unmatched
                    # quote bytes. Printable path extraction remains conservative
                    # and is preferable to rejecting the exact retail revision.
                    flattened = contents
                references = tuple(dict.fromkeys(
                    match.group(0) for match in REFERENCE_PATTERN.finditer(flattened)
                ))
                cubes = tuple(dict.fromkeys(
                    match.group(1)
                    for match in re.finditer(
                        r"(?i)\b(?:camera)?cubemap\s+([A-Za-z0-9_./\\-]+)",
                        flattened,
                    )
                ))
                cached = (references, cubes)
                self._reference_cache[identity] = cached
            for reference in cached[0]:
                resolve_reference(reference)
            for cube in cached[1]:
                add_cube(cube)

        for key in sorted(seeds):
            if pathlib.PurePosixPath(key).suffix in self.INITIAL_TRAVERSABLE:
                pending_files.append(key)
        while pending_files or pending_declarations:
            if pending_files:
                key = pending_files.popleft()
                if key in traversed:
                    continue
                traversed.add(key)
                scan(self.read_text(key), ("file", key))
                continue
            declaration = pending_declarations.popleft()
            if declaration in activated:
                continue
            activated.add(declaration)
            add_key(declaration.key)
            for support_key in declaration.support_keys:
                add_key(support_key)
            scan(declaration.body, declaration)
        return result


def resolve_profile_file(base: pathlib.Path, raw: object) -> pathlib.Path:
    if not isinstance(raw, str) or not raw:
        raise BuildError("profile trace path must be a non-empty string")
    resolved = (base / raw).resolve()
    try:
        resolved.relative_to(base.resolve())
    except ValueError as error:
        raise BuildError(f"profile trace escapes its directory: {raw}") from error
    if not resolved.is_file() or resolved.is_symlink():
        raise BuildError(f"profile trace is missing or not a regular file: {resolved}")
    return resolved


def trace_keys(file: pathlib.Path, approved_archives: set[str], winners: dict[str, Winner]) -> set[str]:
    result: set[str] = set()
    debug_lines = 0
    for raw_line in file.read_text(encoding="utf-8", errors="replace").splitlines():
        match = FS_DEBUG_PATTERN.search(raw_line)
        if match:
            debug_lines += 1
            archive_name = pathlib.PurePath(match.group(2).replace("\\", "/")).name.lower()
            if archive_name not in approved_archives:
                continue
            key, _ = normalize_member(match.group(1).strip())
            winner = winners.get(key)
            if winner is None:
                raise BuildError(f"trace references a missing approved retail member: {match.group(1)}")
            if winner.archive_name != archive_name:
                raise BuildError(
                    f"trace/override disagreement for {winner.name}: traced {archive_name}, "
                    f"winner is {winner.archive_name}"
                )
            result.add(key)
            continue
        if file.suffix == ".assets":
            stripped = raw_line.strip()
            if not stripped or stripped.startswith("#"):
                continue
            key, _ = normalize_member(stripped)
            if key in winners:
                result.add(key)
    if file.suffix != ".assets" and debug_lines == 0:
        raise BuildError(f"profile contains no fs_debug OpenFileRead records: {file}")
    return result


def load_profiles(profile_file: pathlib.Path, retail: dict, winners: dict[str, Winner]) -> tuple[dict, set[str], list[dict]]:
    profile = read_json(profile_file)
    if profile.get("schema") != PROFILE_SCHEMA or profile.get("policy") != PROFILE_POLICY:
        raise BuildError("profile manifest has the wrong schema or policy")
    if profile.get("sourceRevision") != retail["revision"] or profile.get("sourceFingerprint") != retail["fingerprint"]:
        raise BuildError("profile manifest belongs to a different retail revision")
    if profile.get("performancePreset") != PERFORMANCE_PRESET or profile.get("maxMachineSpec") != MAX_MACHINE_SPEC:
        raise BuildError("profiles must be captured with balanced/machineSpec 2")
    coverage = profile.get("coverageScenarios")
    if not isinstance(coverage, list) or not REQUIRED_COVERAGE.issubset(set(coverage)):
        missing = sorted(REQUIRED_COVERAGE - set(coverage or []))
        raise BuildError(f"profile manifest is missing multiplayer coverage: {', '.join(missing)}")
    approved = set(EXPECTED_RETAIL)
    base = profile_file.parent
    common: set[str] = set()
    for raw in profile.get("commonTraces", []):
        common.update(trace_keys(resolve_profile_file(base, raw), approved, winners))
    raw_maps = profile.get("maps")
    if not isinstance(raw_maps, list) or not raw_maps:
        raise BuildError("profile manifest must advertise at least one completed map")
    maps: list[dict] = []
    ids: set[str] = set()
    for raw_map in raw_maps:
        if not isinstance(raw_map, dict):
            raise BuildError("profile map entry must be an object")
        map_id = str(raw_map.get("id", ""))
        if not MAP_PATTERN.fullmatch(map_id) or map_id in ids:
            raise BuildError(f"invalid or duplicate profile map id: {map_id!r}")
        ids.add(map_id)
        traces = raw_map.get("traces")
        if not isinstance(traces, list) or not traces:
            raise BuildError(f"profile map has no traces: {map_id}")
        keys: set[str] = set()
        for raw in traces:
            keys.update(trace_keys(resolve_profile_file(base, raw), approved, winners))
        retail_geometry = raw_map.get("retailGeometry", True)
        if not isinstance(retail_geometry, bool):
            raise BuildError(f"retailGeometry must be boolean: {map_id}")
        if not keys and retail_geometry:
            raise BuildError(f"retail map profile has no approved retail reads: {map_id}")
        if retail_geometry:
            # The render and collision files do not contain the complete entity
            # graph. Requiring the source map lets static closure discover
            # off-screen models, GUIs, effects, and sound shaders.
            for extension in ("cm", "map", "proc"):
                required = f"maps/{map_id}.{extension}".casefold()
                if required not in keys:
                    raise BuildError(f"profile lacks required {extension} geometry for {map_id}")
        maps.append({"id": map_id, "keys": keys, "retailGeometry": retail_geometry})
    closure = RetailDependencyClosure(winners)
    try:
        common = closure.close(common)
        maps = [
            {
                **entry,
                "keys": closure.close(entry["keys"]),
            }
            for entry in maps
        ]
    finally:
        closure.release()
    return profile, common, maps


def partition(common: set[str], maps: list[dict]) -> tuple[set[str], dict[str, set[str]]]:
    users: dict[str, int] = {}
    for entry in maps:
        for key in entry["keys"]:
            users[key] = users.get(key, 0) + 1
    shared = set(common)
    shared.update(key for key, count in users.items() if count >= 2)
    per_map = {entry["id"]: set(entry["keys"]) - shared for entry in maps}
    ownership: dict[str, str] = {key: "common" for key in shared}
    for map_id, keys in per_map.items():
        for key in keys:
            previous = ownership.setdefault(key, map_id)
            if previous != map_id:
                raise BuildError(f"bundle partition conflict for {key}: {previous} and {map_id}")
    return shared, per_map


def pack_directories(keys: set[str], winners: dict[str, Winner]) -> list[str]:
    """Return every parent directory required by idTech's PK4 folder scanner."""
    directories: dict[str, str] = {}
    for key in sorted(keys):
        parts = winners[key].name.split("/")[:-1]
        for length in range(1, len(parts) + 1):
            name = "/".join(parts[:length]) + "/"
            directories.setdefault(name.casefold(), name)
    return sorted(directories.values(), key=lambda name: (name.casefold(), name))


def write_pack(output: pathlib.Path, keys: set[str], winners: dict[str, Winner]) -> None:
    output.parent.mkdir(parents=True, exist_ok=True)
    with zipfile.ZipFile(output, "w", compression=zipfile.ZIP_DEFLATED, compresslevel=9, allowZip64=True) as target:
        sources: dict[pathlib.Path, zipfile.ZipFile] = {}
        try:
            # idTech 4 discovers recursive decl/material/effect folders through
            # explicit ZIP directory entries. Python's ZipFile does not create
            # them when individual members are streamed, so emit the complete
            # parent tree before writing file data.
            for directory in pack_directories(keys, winners):
                info = zipfile.ZipInfo(directory, date_time=FIXED_TIME)
                info.compress_type = zipfile.ZIP_STORED
                info.create_system = 3
                info.external_attr = (0o40755 << 16) | 0x10
                target.writestr(info, b"")
            for key in sorted(keys):
                winner = winners[key]
                source = sources.get(winner.archive_path)
                if source is None:
                    source = zipfile.ZipFile(winner.archive_path)
                    sources[winner.archive_path] = source
                source_info = source.getinfo(winner.member_name)
                info = zipfile.ZipInfo(winner.name, date_time=FIXED_TIME)
                info.compress_type = zipfile.ZIP_DEFLATED
                info.create_system = 3
                info.external_attr = 0o100644 << 16
                info.file_size = source_info.file_size
                with source.open(source_info) as reader, target.open(info, "w") as writer:
                    while chunk := reader.read(CHUNK_SIZE):
                        writer.write(chunk)
        finally:
            for source in sources.values():
                source.close()


def bundle_record(bundle_id: str, kind: str, name: str, file: pathlib.Path) -> dict:
    digest = sha256_file(file)
    return {
        "id": bundle_id,
        "name": name,
        "size": file.stat().st_size,
        "sha256": digest,
        "url": f"/retail/bundles/{digest}-{name}",
        "kind": kind,
    }


def manifest_fingerprint(source_revision: str, source_fingerprint: str, bundles: list[dict], maps: list[dict]) -> str:
    digest = hashlib.sha256()
    digest.update(f"{source_revision}\0{source_fingerprint}\n".encode())
    for bundle in sorted(bundles, key=lambda entry: entry["id"]):
        digest.update(
            (f"{bundle['id']}\0{bundle['kind']}\0{bundle['name']}\0{bundle['size']}\0"
             f"{bundle['sha256']}\0{bundle['url']}\n").encode()
        )
    for entry in sorted(maps, key=lambda item: item["id"]):
        digest.update(f"{entry['id']}\0{','.join(entry['bundleIds'])}\n".encode())
    return digest.hexdigest()


def extra_field_ids(extra: bytes, location: str) -> set[int]:
    ids: set[int] = set()
    offset = 0
    while offset < len(extra):
        if offset + 4 > len(extra):
            raise BuildError(f"truncated ZIP extra field in {location}")
        field_id, size = struct.unpack_from("<HH", extra, offset)
        offset += 4
        if offset + size > len(extra):
            raise BuildError(f"truncated ZIP extra payload in {location}")
        ids.add(field_id)
        offset += size
    return ids


def verify_standard_pk4_entry(pack: zipfile.ZipFile, member: zipfile.ZipInfo, compression: int) -> None:
    if member.extract_version > 20 or member.compress_type != compression:
        raise BuildError(f"bundle member requires unsupported ZIP features: {member.filename}")
    if 0x0001 in extra_field_ids(member.extra, f"central header for {member.filename}"):
        raise BuildError(f"bundle member has a ZIP64 central extra field: {member.filename}")
    if pack.fp is None:
        raise BuildError(f"bundle closed while checking local header: {member.filename}")
    pack.fp.seek(member.header_offset)
    header = pack.fp.read(30)
    if len(header) != 30 or struct.unpack_from("<I", header)[0] != 0x04034B50:
        raise BuildError(f"invalid local ZIP header: {member.filename}")
    extract_version, compression = struct.unpack_from("<HxxH", header, 4)
    name_size, extra_size = struct.unpack_from("<HH", header, 26)
    local_name = pack.fp.read(name_size)
    local_extra = pack.fp.read(extra_size)
    if local_name != member.filename.encode("utf-8"):
        raise BuildError(f"local ZIP member name mismatch: {member.filename}")
    if extract_version > 20 or compression != member.compress_type:
        raise BuildError(f"bundle local header requires unsupported ZIP features: {member.filename}")
    if 0x0001 in extra_field_ids(local_extra, f"local header for {member.filename}"):
        raise BuildError(f"bundle member has a ZIP64 local extra field: {member.filename}")


def verify_pack(file: pathlib.Path, expected_keys: set[str] | None, winners: dict[str, Winner]) -> set[str]:
    try:
        with zipfile.ZipFile(file) as pack:
            seen: set[str] = set()
            seen_directories: dict[str, str] = {}
            for member in pack.infolist():
                if member.is_dir():
                    verify_standard_pk4_entry(pack, member, zipfile.ZIP_STORED)
                    if member.file_size != 0 or member.CRC != 0:
                        raise BuildError(f"bundle directory contains data: {member.filename}")
                    raw = member.filename[:-1]
                    key, name = normalize_member(raw)
                    directory_name = name + "/"
                    if key in seen_directories:
                        raise BuildError(f"bundle has a duplicate directory: {directory_name}")
                    seen_directories[key] = directory_name
                    continue
                verify_standard_pk4_entry(pack, member, zipfile.ZIP_DEFLATED)
                key, name = normalize_member(member.filename)
                if key in seen:
                    raise BuildError(f"bundle has a duplicate member: {name}")
                seen.add(key)
                expected = winners.get(key)
                if expected is None or expected.name != name or expected.size != member.file_size or expected.crc != member.CRC:
                    raise BuildError(f"bundle member differs from retail winner: {name}")
                with pack.open(member) as stream:
                    while stream.read(CHUNK_SIZE):
                        pass
            if expected_keys is not None and seen != expected_keys:
                raise BuildError(f"bundle member set mismatch: {file.name}")
            expected_directories = {
                name[:-1].casefold(): name for name in pack_directories(seen, winners)
            }
            if seen_directories != expected_directories:
                raise BuildError(f"bundle directory tree mismatch: {file.name}")
            return seen
    except zipfile.BadZipFile as error:
        raise BuildError(f"invalid generated bundle {file}: {error}") from error


def stage(args: argparse.Namespace) -> dict:
    retail, files = load_retail(args.retail_dir.resolve(), args.retail_manifest.resolve())
    winners = build_override_graph(files)
    profile, explicit_common, maps = load_profiles(args.profiles.resolve(), retail, winners)
    common, per_map = partition(explicit_common, maps)
    if not common:
        raise BuildError("common bundle would be empty; capture both q4web and a retail map in fresh processes")
    retail_geometry_by_id = {entry["id"]: entry["retailGeometry"] for entry in maps}
    for map_id, keys in per_map.items():
        if not keys and retail_geometry_by_id[map_id]:
            raise BuildError(f"map bundle would be empty: {map_id}")
    output = args.output.resolve()
    output.mkdir(parents=True, exist_ok=True)
    temporary = pathlib.Path(tempfile.mkdtemp(prefix=".retail-bundles-", dir=output))
    records: list[dict] = []
    expected_by_name: dict[str, set[str]] = {}
    try:
        # OpenQ4 only treats q4base as a valid game directory when the
        # canonical retail sentinel exists at q4base/pak001.pk4.
        common_name = "pak001.pk4"
        write_pack(temporary / common_name, common, winners)
        expected_by_name[common_name] = common
        records.append(bundle_record("common", "common", common_name, temporary / common_name))
        map_records: list[dict] = []
        for entry in sorted(maps, key=lambda item: item["id"]):
            slug = entry["id"].replace("/", "_")
            name = f"pak910_q4retail_{slug}.pk4"
            write_pack(temporary / name, per_map[entry["id"]], winners)
            expected_by_name[name] = per_map[entry["id"]]
            bundle_id = f"map/{entry['id']}"
            records.append(bundle_record(bundle_id, "map", name, temporary / name))
            map_records.append({"id": entry["id"], "bundleIds": ["common", bundle_id]})
        for record in records:
            verify_pack(temporary / record["name"], expected_by_name[record["name"]], winners)
        fingerprint = manifest_fingerprint(retail["revision"], retail["fingerprint"], records, map_records)
        manifest = {
            "schema": BUNDLE_SCHEMA,
            "policy": BUNDLE_POLICY,
            "revision": f"q4bundles-{fingerprint[:20]}",
            "fingerprint": fingerprint,
            "sourceRevision": retail["revision"],
            "sourceFingerprint": retail["fingerprint"],
            "performancePreset": PERFORMANCE_PRESET,
            "maxMachineSpec": MAX_MACHINE_SPEC,
            "coverage": str(profile.get("coverage", "beta")),
            "coverageScenarios": list(profile["coverageScenarios"]),
            "totalBytes": sum(record["size"] for record in records),
            "bundles": records,
            "maps": map_records,
        }
        encoded = json.dumps(manifest, indent=2, ensure_ascii=True) + "\n"
        temporary_manifest = temporary / "manifest.json"
        temporary_manifest.write_text(encoded, encoding="utf-8")
        os.chmod(temporary_manifest, 0o644)
        keep = {record["name"] for record in records}
        for name in sorted(keep):
            os.replace(temporary / name, output / name)
        for old in output.glob("pak9*_q4retail_*.pk4"):
            if old.name not in keep:
                old.unlink()
        manifest_file = args.manifest.resolve()
        manifest_file.parent.mkdir(parents=True, exist_ok=True)
        os.replace(temporary_manifest, manifest_file)
        return manifest
    finally:
        shutil.rmtree(temporary, ignore_errors=True)


def verify_staged(args: argparse.Namespace) -> dict:
    retail, files = load_retail(args.retail_dir.resolve(), args.retail_manifest.resolve())
    winners = build_override_graph(files)
    manifest = read_json(args.manifest.resolve())
    if manifest.get("schema") != BUNDLE_SCHEMA or manifest.get("policy") != BUNDLE_POLICY:
        raise BuildError("bundle manifest has the wrong schema or policy")
    if manifest.get("sourceRevision") != retail["revision"] or manifest.get("sourceFingerprint") != retail["fingerprint"]:
        raise BuildError("bundle manifest belongs to a different retail revision")
    if manifest.get("performancePreset") != PERFORMANCE_PRESET or manifest.get("maxMachineSpec") != MAX_MACHINE_SPEC:
        raise BuildError("bundle manifest has the wrong performance profile")
    bundles = manifest.get("bundles")
    maps = manifest.get("maps")
    if not isinstance(bundles, list) or len(bundles) < 2 or not isinstance(maps, list) or not maps:
        raise BuildError("bundle manifest has no usable bundle graph")
    output = args.output.resolve()
    bundle_ids: set[str] = set()
    bundle_names: set[str] = set()
    kind_by_id: dict[str, str] = {}
    all_keys: set[str] = set()
    total = 0
    for record in bundles:
        if not isinstance(record, dict):
            raise BuildError("invalid bundle record")
        bundle_id = str(record.get("id", ""))
        name = str(record.get("name", ""))
        kind = str(record.get("kind", ""))
        size = record.get("size")
        digest = str(record.get("sha256", ""))
        file = output / name
        if not re.fullmatch(r"[a-z0-9][a-z0-9._/-]{0,119}", bundle_id) or bundle_id in bundle_ids:
            raise BuildError(f"invalid or duplicate bundle id: {bundle_id}")
        if not re.fullmatch(r"[a-z0-9][a-z0-9._-]{0,154}\.pk4", name) or name in bundle_names:
            raise BuildError(f"invalid or duplicate bundle name: {name}")
        if kind not in {"common", "map", "botnav"} or not isinstance(size, int) or size <= 0 or not SHA256_PATTERN.fullmatch(digest):
            raise BuildError(f"invalid bundle metadata: {bundle_id}")
        if record.get("url") != f"/retail/bundles/{digest}-{name}":
            raise BuildError(f"non-canonical bundle URL: {bundle_id}")
        if not file.is_file() or file.is_symlink() or file.stat().st_size != size or sha256_file(file) != digest:
            raise BuildError(f"bundle bytes differ from manifest: {name}")
        keys = verify_pack(file, None, winners)
        overlap = all_keys.intersection(keys)
        if overlap:
            raise BuildError(f"bundle member collision: {sorted(overlap)[0]}")
        all_keys.update(keys)
        bundle_ids.add(bundle_id)
        bundle_names.add(name)
        kind_by_id[bundle_id] = kind
        total += size
    common_bundles = [record for record in bundles if record.get("kind") == "common"]
    if len(common_bundles) != 1 or common_bundles[0].get("name") != "pak001.pk4":
        raise BuildError("bundle manifest must contain exactly one common bundle named pak001.pk4")
    if manifest.get("totalBytes") != total:
        raise BuildError("bundle totalBytes mismatch")
    map_ids: set[str] = set()
    for entry in maps:
        map_id = str(entry.get("id", "")) if isinstance(entry, dict) else ""
        ids = entry.get("bundleIds") if isinstance(entry, dict) else None
        if not MAP_PATTERN.fullmatch(map_id) or map_id in map_ids or not isinstance(ids, list) or len(ids) < 2:
            raise BuildError(f"invalid map bundle graph: {map_id}")
        if len(set(ids)) != len(ids) or any(bundle_id not in bundle_ids for bundle_id in ids):
            raise BuildError(f"invalid map bundle references: {map_id}")
        if not any(kind_by_id[bundle_id] == "common" for bundle_id in ids) or not any(kind_by_id[bundle_id] == "map" for bundle_id in ids):
            raise BuildError(f"map graph lacks common or map bundle: {map_id}")
        map_ids.add(map_id)
    fingerprint = manifest_fingerprint(retail["revision"], retail["fingerprint"], bundles, maps)
    if manifest.get("fingerprint") != fingerprint or manifest.get("revision") != f"q4bundles-{fingerprint[:20]}":
        raise BuildError("bundle manifest fingerprint is not canonical")
    return manifest


def parse_args() -> argparse.Namespace:
    root = pathlib.Path(__file__).resolve().parent.parent
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--retail-dir", type=pathlib.Path, default=root / "assets-src/gog/q4base")
    parser.add_argument("--retail-manifest", type=pathlib.Path, default=root / "runtime/retail-manifest.json")
    parser.add_argument("--profiles", type=pathlib.Path, default=root / "runtime/retail-bundle-profiles/manifest.json")
    parser.add_argument("--output", type=pathlib.Path, default=root / "runtime/retail-bundles")
    parser.add_argument("--manifest", type=pathlib.Path, default=root / "runtime/retail-bundle-manifest.json")
    parser.add_argument("--verify", action="store_true")
    return parser.parse_args()


def main() -> int:
    try:
        args = parse_args()
        manifest = verify_staged(args) if args.verify else stage(args)
    except (BuildError, OSError, zipfile.BadZipFile) as error:
        print(f"retail bundle staging failed: {error}", file=sys.stderr)
        return 1
    print(
        f"{'Verified' if args.verify else 'Retail bundles'} {manifest['revision']}: {len(manifest['maps'])} maps, "
        f"{len(manifest['bundles'])} packs, {manifest['totalBytes']} bytes"
    )
    return 0


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