#!/usr/bin/env python3
"""Reject public PK4 entries that duplicate operator-owned retail data."""

from __future__ import annotations

import argparse
import hashlib
import json
import pathlib
import stat
import sys
import zipfile
from dataclasses import dataclass


CHUNK_SIZE = 1024 * 1024


@dataclass(frozen=True)
class Entry:
    archive: str
    path: str
    sha256: str
    size: int

    @property
    def key(self) -> tuple[str, str, str, str, str]:
        return (self.archive, self.path, self.sha256, "", "")


def normalized_member(name: str) -> str:
    value = name.replace("\\", "/").lstrip("/")
    parts = value.split("/")
    if not value or value != name.replace("\\", "/") or any(
        part in {"", ".", ".."} for part in parts
    ):
        raise ValueError(f"unsafe archive member path: {name!r}")
    return "/".join(parts).casefold()


def digest_member(archive: zipfile.ZipFile, info: zipfile.ZipInfo) -> str:
    digest = hashlib.sha256()
    with archive.open(info, "r") as source:
        while chunk := source.read(CHUNK_SIZE):
            digest.update(chunk)
    return digest.hexdigest()


def read_archive(file: pathlib.Path, root: pathlib.Path) -> list[Entry]:
    result: list[Entry] = []
    names: set[str] = set()
    try:
        with zipfile.ZipFile(file, "r") as archive:
            for info in archive.infolist():
                if info.is_dir():
                    continue
                path = normalized_member(info.filename)
                if path in names:
                    raise ValueError(f"duplicate member path in {file}: {path}")
                names.add(path)
                mode = info.external_attr >> 16
                if stat.S_ISLNK(mode):
                    raise ValueError(f"symlink member in {file}: {path}")
                if info.flag_bits & 0x1:
                    raise ValueError(f"encrypted member in {file}: {path}")
                result.append(
                    Entry(
                        archive=file.relative_to(root).as_posix(),
                        path=path,
                        sha256=digest_member(archive, info),
                        size=info.file_size,
                    )
                )
    except zipfile.BadZipFile as error:
        raise ValueError(f"invalid PK4 archive {file}: {error}") from error
    return result


def archives(directory: pathlib.Path) -> list[pathlib.Path]:
    if not directory.is_dir():
        raise ValueError(f"PK4 directory does not exist: {directory}")
    files = sorted(
        (entry for entry in directory.iterdir() if entry.is_file() and entry.suffix.casefold() == ".pk4"),
        key=lambda entry: entry.name.casefold(),
    )
    if not files:
        raise ValueError(f"PK4 directory contains no archives: {directory}")
    return files


def load_entries(directory: pathlib.Path) -> list[Entry]:
    result: list[Entry] = []
    for file in archives(directory):
        result.extend(read_archive(file, directory))
    return result


def allowlist_key(value: dict[str, object]) -> tuple[str, str, str, str, str]:
    fields = ("publicArchive", "publicPath", "sha256", "retailArchive", "retailPath")
    if any(not isinstance(value.get(field), str) or not value[field] for field in fields):
        raise ValueError("allowlist entry is missing an exact archive/path/SHA tuple")
    if not isinstance(value.get("reason"), str) or not value["reason"].strip():
        raise ValueError("allowlist entry has no reason")
    if not isinstance(value.get("licenseEvidence"), str) or not value["licenseEvidence"].strip():
        raise ValueError("allowlist entry has no license evidence")
    return tuple(str(value[field]).casefold() for field in fields)  # type: ignore[return-value]


def load_allowlist(file: pathlib.Path) -> set[tuple[str, str, str, str, str]]:
    payload = json.loads(file.read_text(encoding="utf-8"))
    if payload.get("schema") != 1 or not isinstance(payload.get("entries"), list):
        raise ValueError(f"invalid duplicate allowlist: {file}")
    keys = [allowlist_key(entry) for entry in payload["entries"] if isinstance(entry, dict)]
    if len(keys) != len(payload["entries"]) or len(set(keys)) != len(keys):
        raise ValueError("duplicate or malformed allowlist entry")
    return set(keys)


def collision_key(public: Entry, retail: Entry) -> tuple[str, str, str, str, str]:
    return (
        public.archive.casefold(),
        public.path.casefold(),
        public.sha256,
        retail.archive.casefold(),
        retail.path.casefold(),
    )


def audit(public_dir: pathlib.Path, retail_dirs: list[pathlib.Path], allowlist_file: pathlib.Path) -> dict[str, object]:
    public = load_entries(public_dir)
    retail: list[Entry] = []
    for directory in retail_dirs:
        retail.extend(load_entries(directory))

    by_digest: dict[str, list[Entry]] = {}
    for entry in retail:
        by_digest.setdefault(entry.sha256, []).append(entry)

    allowed = load_allowlist(allowlist_file)
    used: set[tuple[str, str, str, str, str]] = set()
    blocked: list[tuple[Entry, Entry]] = []
    for candidate in public:
        for private in by_digest.get(candidate.sha256, []):
            key = collision_key(candidate, private)
            if key in allowed:
                used.add(key)
            else:
                blocked.append((candidate, private))

    stale = sorted(allowed - used)
    if stale:
        raise ValueError(f"stale duplicate allowlist entries: {json.dumps(stale)}")
    if blocked:
        lines = ["public PK4 contains exact retail duplicates:"]
        for candidate, private in blocked[:50]:
            relation = "same-path" if candidate.path == private.path else "same-bytes"
            lines.append(
                f"  {candidate.archive}:{candidate.path} == "
                f"{private.archive}:{private.path} ({relation}, sha256={candidate.sha256})"
            )
        if len(blocked) > 50:
            lines.append(f"  ... and {len(blocked) - 50} more")
        raise ValueError("\n".join(lines))

    return {
        "ok": True,
        "publicArchives": len(archives(public_dir)),
        "publicEntries": len(public),
        "retailArchives": sum(len(archives(directory)) for directory in retail_dirs),
        "retailEntries": len(retail),
        "allowlistEntries": len(allowed),
        "duplicates": 0,
    }


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--public-dir", required=True, type=pathlib.Path)
    parser.add_argument("--retail-dir", required=True, action="append", type=pathlib.Path)
    parser.add_argument("--allowlist", required=True, type=pathlib.Path)
    parser.add_argument("--report", type=pathlib.Path)
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    try:
        report = audit(args.public_dir.resolve(), [item.resolve() for item in args.retail_dir], args.allowlist.resolve())
    except (OSError, ValueError, KeyError, json.JSONDecodeError) as error:
        print(f"content audit failed: {error}", file=sys.stderr)
        return 1
    encoded = json.dumps(report, indent=2, sort_keys=True) + "\n"
    if args.report:
        args.report.parent.mkdir(parents=True, exist_ok=True)
        args.report.write_text(encoded, encoding="utf-8")
    sys.stdout.write(encoded)
    return 0


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