#!/usr/bin/env python3
"""Fail-closed license, symbol, architecture and slice audit for VLC modules."""

from __future__ import annotations

import argparse
import hashlib
import re
import subprocess
import sys
from dataclasses import dataclass
from pathlib import Path

from macho_audit import audit_archive_identity, run_xcrun

GPL_MARKER = b"Licensed under the terms of the GNU General Public License, version 2 or later."
LGPL_MARKER = b"Licensed under the terms of the GNU Lesser General Public License, version 2.1 or later."
ARCHIVE_RE = re.compile(r"^lib[A-Za-z0-9_+-]+_plugin[.]a$")
ARCH_RE = re.compile(r"^[A-Za-z0-9_+-]+$")
ENTRY_RE = re.compile(r"^_vlc_entry__[A-Za-z0-9_]+$")
LICENSE_ENTRY_RE = re.compile(r"^_vlc_entry_license__[A-Za-z0-9_]+$")
FORBIDDEN_MODULE_ARCHIVE_TOKENS = ("zvbi",)


@dataclass(frozen=True)
class Slice:
    kind: str
    architecture: str
    directory: Path


@dataclass(frozen=True)
class Inspection:
    sha256: str
    classification: str
    entry: str
    license_entry: str


def sha256_bytes(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()


def digest_rows(rows: list[str]) -> str:
    return hashlib.sha256("".join(f"{row}\n" for row in sorted(rows)).encode()).hexdigest()


def count_exact_c_strings(data: bytes, marker: bytes) -> int:
    needle = marker + b"\0"
    offset = 0
    count = 0
    while True:
        position = data.find(needle, offset)
        if position < 0:
            return count
        if position == 0 or data[position - 1] == 0:
            count += 1
        offset = position + 1


def contains_exact_c_string(data: bytes, marker: bytes) -> bool:
    return count_exact_c_strings(data, marker) > 0


def classify(data: bytes) -> str:
    gpl_count = count_exact_c_strings(data, GPL_MARKER)
    lgpl_count = count_exact_c_strings(data, LGPL_MARKER)
    if gpl_count == 1 and lgpl_count == 0:
        return "gpl"
    if lgpl_count == 1 and gpl_count == 0:
        return "lgpl"
    if gpl_count and lgpl_count:
        return "mixed"
    if gpl_count > 1 or lgpl_count > 1:
        return "duplicate-marker"
    return "unknown"


def defined_module_symbols(path: Path) -> tuple[str, str]:
    symbols = [
        line.strip()
        for line in run_xcrun("nm", "-gUj", str(path)).splitlines()
        if line.strip().startswith("_vlc_entry")
    ]
    entries = sorted(symbol for symbol in symbols if ENTRY_RE.fullmatch(symbol))
    licenses = sorted(symbol for symbol in symbols if LICENSE_ENTRY_RE.fullmatch(symbol))
    if len(entries) != 1 or len(licenses) != 1:
        raise ValueError(
            f"{path} must define exactly one module entry and license entry; "
            f"found entries={entries}, licenses={licenses}"
        )
    expected_license = entries[0].replace("_vlc_entry__", "_vlc_entry_license__", 1)
    if licenses[0] != expected_license:
        raise ValueError(
            f"{path} entry/license identities differ: {entries[0]} != {licenses[0]}"
        )
    return entries[0], licenses[0]


def inspect_archive(
    path: Path, item: Slice, os_style: str, minimum_os: str, forbidden_prefix: str
) -> Inspection:
    if path.is_symlink() or not path.is_file() or not ARCHIVE_RE.fullmatch(path.name):
        raise ValueError(f"invalid static module archive: {path}")
    data = path.read_bytes()
    if forbidden_prefix.encode() in data or b"/Users/" in data:
        raise ValueError(f"unmapped build path found in static module archive: {path}")
    classification = classify(data)
    if classification not in {"gpl", "lgpl"}:
        raise ValueError(f"{path} has an unreviewable {classification} license marker")
    audit_archive_identity(
        path,
        architecture=item.architecture,
        os_style=os_style,
        kind=item.kind,
        minimum_os=minimum_os,
    )
    entry, license_entry = defined_module_symbols(path)
    return Inspection(sha256_bytes(data), classification, entry, license_entry)


def parse_arches(value: str) -> list[str]:
    arches = value.split()
    for arch in arches:
        if not ARCH_RE.fullmatch(arch):
            raise ValueError(f"invalid architecture: {arch!r}")
    return arches


def actual_arch(arch: str) -> str:
    return "arm64" if arch == "aarch64" else arch


def slices_for(
    vlc_root: Path, os_style: str, device_arches: list[str], simulator_arches: list[str]
) -> list[Slice]:
    slices: list[Slice] = []
    for kind, arches in (("device", device_arches), ("simulator", simulator_arches)):
        suffix = "OS" if kind == "device" else "Simulator"
        for arch in arches:
            slices.append(
                Slice(
                    kind,
                    actual_arch(arch),
                    vlc_root
                    / f"install-{os_style}{suffix}"
                    / actual_arch(arch)
                    / "lib/vlc/plugins",
                )
            )
    return slices


def read_assignments(path: Path) -> dict[str, str]:
    values: dict[str, str] = {}
    for raw in path.read_text(encoding="utf-8").splitlines():
        line = raw.strip()
        if not line or line.startswith("#"):
            continue
        key, separator, value = line.partition("=")
        if not separator or not re.fullmatch(r"[A-Z][A-Z0-9_]*", key):
            raise ValueError(f"invalid assignment in {path}: {raw}")
        values[key] = value
    return values


def baseline_names(path: Path) -> set[str]:
    names = {
        line.strip()
        for line in path.read_text(encoding="utf-8").splitlines()
        if line.strip() and not line.lstrip().startswith("#")
    }
    invalid = sorted(name for name in names if not ARCHIVE_RE.fullmatch(name))
    if invalid:
        raise ValueError(f"invalid baseline module names: {invalid}")
    return names


def write_lines(path: Path, header: str, rows: list[str]) -> None:
    path.write_text(
        header + "\n" + "".join(f"{row}\n" for row in sorted(rows)), encoding="utf-8"
    )


def reject_symlink_components(path: Path, anchor: Path) -> None:
    relative = path.relative_to(anchor)
    current = anchor
    if current.is_symlink():
        raise ValueError(f"symlinked audit root is forbidden: {current}")
    for part in relative.parts:
        current = current / part
        if current.is_symlink():
            raise ValueError(f"symlinked path component is forbidden: {current}")


def reject_absolute_symlink_components(path: Path) -> None:
    absolute = path.absolute()
    current = Path(absolute.anchor)
    for part in absolute.parts[1:]:
        current = current / part
        if current.is_symlink():
            raise ValueError(f"symlinked absolute path component is forbidden: {current}")


def audit(
    vlc_root: Path,
    os_style: str,
    device_arches: list[str],
    simulator_arches: list[str],
    audit_dir: Path,
    minimum_os: str,
    pins: dict[str, str],
    required_baseline: set[str],
    forbidden_prefix: str,
) -> list[str]:
    slices = slices_for(vlc_root, os_style, device_arches, simulator_arches)
    reject_absolute_symlink_components(vlc_root)
    expected_arches = {"device": {"arm64"}, "simulator": {"arm64", "x86_64"}}
    actual_arches = {
        kind: {item.architecture for item in slices if item.kind == kind}
        for kind in expected_arches
    }
    if actual_arches != expected_arches:
        raise ValueError(f"unexpected candidate slice architecture sets: {actual_arches}")

    relative_sets: dict[Slice, set[str]] = {}
    inspections: dict[Slice, dict[str, Inspection]] = {}
    inventory_rows: list[str] = []
    included_rows: list[str] = []
    excluded_rows: list[str] = []

    for item in slices:
        reject_symlink_components(item.directory, vlc_root)
        if item.directory.is_symlink() or not item.directory.is_dir():
            raise ValueError(f"missing or symlinked static module directory: {item.directory}")
        paths = sorted(item.directory.rglob("*.a"))
        if not paths:
            raise ValueError(f"no static modules found in {item.directory}")
        relative = {path.relative_to(item.directory).as_posix() for path in paths}
        if len(relative) != len(paths):
            raise ValueError(f"duplicate static module paths beneath {item.directory}")
        relative_sets[item] = relative
        inspections[item] = {}
        for path in paths:
            reject_symlink_components(path, item.directory)
            relative_name = path.relative_to(item.directory).as_posix()
            inspected = inspect_archive(path, item, os_style, minimum_os, forbidden_prefix)
            inspections[item][relative_name] = inspected
            row = (
                f"{item.kind}\t{item.architecture}\t{relative_name}\t{inspected.sha256}"
                f"\t{inspected.classification}\t{inspected.entry}\t{inspected.license_entry}"
            )
            inventory_rows.append(row)
            (included_rows if inspected.classification == "lgpl" else excluded_rows).append(row)

        by_basename: dict[str, list[str]] = {}
        for relative_name in relative:
            by_basename.setdefault(Path(relative_name).name, []).append(relative_name)
        for basename, copies in sorted(by_basename.items()):
            root = [name for name in copies if name == basename]
            category = [name for name in copies if len(Path(name).parts) == 2]
            if len(copies) != 2 or len(root) != 1 or len(category) != 1:
                raise ValueError(
                    f"{item.kind}/{item.architecture}/{basename} must have exactly one "
                    f"flat and one category copy; found {copies}"
                )
            identities = {inspections[item][name] for name in copies}
            if len(identities) != 1:
                raise ValueError(
                    f"flat/category copies differ for {item.kind}/{item.architecture}/{basename}"
                )

    reference = slices[0]
    for item in slices[1:]:
        if relative_sets[item] != relative_sets[reference]:
            raise ValueError(f"full module path inventory differs for {item.kind}/{item.architecture}")

    root_names = sorted(name for name in relative_sets[reference] if "/" not in name)
    semantic: dict[str, Inspection] = {
        name: inspections[reference][name] for name in root_names
    }
    forbidden_modules = sorted(
        name
        for name in semantic
        if any(token in name.lower() for token in FORBIDDEN_MODULE_ARCHIVE_TOKENS)
    )
    if forbidden_modules:
        raise ValueError(f"forbidden VLC modules found: {forbidden_modules}")
    for item in slices[1:]:
        for name in root_names:
            current = inspections[item][name]
            expected = semantic[name]
            if (current.classification, current.entry, current.license_entry) != (
                expected.classification,
                expected.entry,
                expected.license_entry,
            ):
                raise ValueError(
                    f"module identity differs for {item.kind}/{item.architecture}/{name}"
                )

    approved_rows = [
        f"{name}\t{item.entry.removeprefix('_')}\t{item.license_entry.removeprefix('_')}"
        f"\t{item.classification}"
        for name, item in sorted(semantic.items())
    ]
    included = sorted(name for name, item in semantic.items() if item.classification == "lgpl")
    excluded = sorted(name for name, item in semantic.items() if item.classification == "gpl")
    missing_baseline = sorted(required_baseline - set(included))
    if missing_baseline:
        raise ValueError(f"required LGPL playback modules are missing: {missing_baseline}")

    prefix = "BEAM_IPHONE" if os_style == "iPhone" else "BEAM_APPLETV"
    digest_inputs = {
        f"{prefix}_MODULE_INVENTORY_SHA256": approved_rows,
        f"{prefix}_MODULE_INCLUDED_SHA256": [
            row for row in approved_rows if row.endswith("\tlgpl")
        ],
        f"{prefix}_MODULE_EXCLUDED_SHA256": [
            row for row in approved_rows if row.endswith("\tgpl")
        ],
    }
    digest_report: list[str] = []
    digest_mismatches: list[str] = []
    for key, rows in digest_inputs.items():
        actual = digest_rows(rows)
        expected = pins.get(key, "")
        if actual != expected:
            digest_mismatches.append(f"{key}: calculated {actual}, reviewed {expected}")
        digest_report.append(f"{key}={actual}")

    audit_dir.mkdir(parents=True, exist_ok=True)
    header = "slice\tarchitecture\tarchive\tsha256\tclassification\tentry\tlicense_entry"
    write_lines(audit_dir / "module-inventory.tsv", header, inventory_rows)
    write_lines(audit_dir / "included-lgpl-modules.tsv", header, included_rows)
    write_lines(audit_dir / "excluded-gpl-modules.tsv", header, excluded_rows)
    write_lines(
        audit_dir / "approved-module-entries.tsv",
        "archive\tentry\tlicense_entry\tclassification",
        approved_rows,
    )
    (audit_dir / "module-inventory-digests.txt").write_text(
        "\n".join(digest_report) + "\n", encoding="utf-8"
    )
    if digest_mismatches:
        raise ValueError(
            "module inventory is not the reviewed pinned set; audit evidence was written:\n  "
            + "\n  ".join(digest_mismatches)
        )
    return included


def parse_args(argv: list[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--vlc-root", required=True, type=Path)
    parser.add_argument("--os-style", required=True, choices=("iPhone", "AppleTV"))
    parser.add_argument("--device-arches", required=True)
    parser.add_argument("--simulator-arches", required=True)
    parser.add_argument("--audit-dir", required=True, type=Path)
    parser.add_argument("--minimum-os", required=True)
    parser.add_argument("--pins", required=True, type=Path)
    parser.add_argument("--baseline", required=True, type=Path)
    parser.add_argument("--forbidden-prefix", required=True)
    return parser.parse_args(argv)


def main(argv: list[str] | None = None) -> int:
    args = parse_args(sys.argv[1:] if argv is None else argv)
    try:
        included = audit(
            args.vlc_root.absolute(),
            args.os_style,
            parse_arches(args.device_arches),
            parse_arches(args.simulator_arches),
            args.audit_dir.resolve(),
            args.minimum_os,
            read_assignments(args.pins),
            baseline_names(args.baseline),
            args.forbidden_prefix,
        )
    except (OSError, ValueError, subprocess.CalledProcessError) as error:
        print(f"ERROR: {error}", file=sys.stderr)
        return 1
    # The patched upstream script captures stdout as its linker module list.
    print("\n".join(included))
    return 0


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