#!/usr/bin/env python3
"""Read-only, opaque structural feature probes for extracted disc files.

The probe validates a manifest-tracked extraction root before reading fixed-size
prefixes. Successful reports use deterministic opaque candidate identifiers and
never emit source paths or raw source bytes. It intentionally reports heuristics,
not file-format conclusions.
"""

from __future__ import annotations

import argparse
from collections import Counter
from dataclasses import dataclass
import hashlib
import math
from pathlib import Path
import re
import sys
from typing import Sequence


MANIFEST_HEADER = "path\tlba\textents\tsize\tsha1"
SAFE_COMPONENT = re.compile(r"[A-Za-z0-9._;-]+\Z")
ISO_VERSION_SUFFIX = re.compile(r";[0-9]+\Z")
SHA1_PATTERN = re.compile(r"[0-9a-f]{40}\Z")
MAX_PREFIX_SIZE = 1024 * 1024


class ProbeValidationError(ValueError):
    """Raised when probe input cannot safely support a structural survey."""


@dataclass(frozen=True)
class ManifestEntry:
    """One validated manifest row, retained only for local probe selection."""

    path: str
    lba: int
    size: int
    sha1: str


@dataclass(frozen=True)
class CandidateFeatures:
    """Non-content structural measurements for one opaque candidate."""

    identifier: str
    size: int
    prefix_size: int
    prefix_sha1: str
    entropy_milli: int
    zero_bytes: int
    monotonic_u32_run_count: int


@dataclass(frozen=True)
class PrefixFeatureReport:
    """A deterministic, opaque fixed-prefix feature report."""

    extension: str
    candidate_count: int
    common_prefix_span: int
    constant_positions: int
    candidates: tuple[CandidateFeatures, ...]


def _validate_manifest_path(path: str, context: str) -> None:
    if not path:
        raise ProbeValidationError(f"{context}: empty path")
    for component in path.split("/"):
        if not component:
            raise ProbeValidationError(f"{context}: empty path component")
        if component in {".", ".."}:
            raise ProbeValidationError(f"{context}: unsafe path component")
        try:
            component.encode("ascii")
        except UnicodeEncodeError as exc:
            raise ProbeValidationError(f"{context}: non-ASCII path component") from exc
        if SAFE_COMPONENT.fullmatch(component) is None:
            raise ProbeValidationError(f"{context}: unsafe path component")


def _parse_extents(field: str, context: str) -> tuple[tuple[int, int], ...]:
    if not field:
        raise ProbeValidationError(f"{context}: empty extent field")
    extents: list[tuple[int, int]] = []
    for encoded_extent in field.split(","):
        lba_text, separator, size_text = encoded_extent.partition(":")
        if not separator or ":" in size_text:
            raise ProbeValidationError(f"{context}: malformed extent")
        if not lba_text.isdecimal() or not size_text.isdecimal():
            raise ProbeValidationError(f"{context}: non-decimal extent")
        extents.append((int(lba_text), int(size_text)))
    return tuple(extents)


def read_manifest(manifest_path: str | Path) -> tuple[ManifestEntry, ...]:
    """Read a strict extraction manifest without returning hashes or source data."""
    manifest = Path(manifest_path)
    if manifest.is_symlink():
        raise ProbeValidationError("manifest must not be a symbolic link")
    try:
        text = manifest.read_text(encoding="ascii")
    except UnicodeDecodeError as exc:
        raise ProbeValidationError("manifest is not ASCII") from exc
    except OSError as exc:
        raise ProbeValidationError(f"cannot read manifest: {exc}") from exc

    lines = text.splitlines()
    if not lines or lines[0] != MANIFEST_HEADER:
        raise ProbeValidationError("manifest has an unexpected header")

    entries: list[ManifestEntry] = []
    seen_paths: set[str] = set()
    for line_number, line in enumerate(lines[1:], start=2):
        fields = line.split("\t")
        if len(fields) != 5:
            raise ProbeValidationError(f"manifest line {line_number}: expected five fields")
        path, lba_text, extent_text, size_text, sha1 = fields
        context = f"manifest line {line_number}"
        _validate_manifest_path(path, context)
        if path in seen_paths:
            raise ProbeValidationError(f"{context}: duplicate path")
        seen_paths.add(path)
        if not lba_text.isdecimal() or not size_text.isdecimal():
            raise ProbeValidationError(f"{context}: non-decimal LBA or size")
        if SHA1_PATTERN.fullmatch(sha1) is None:
            raise ProbeValidationError(f"{context}: malformed SHA-1")
        extents = _parse_extents(extent_text, context)
        lba = int(lba_text)
        size = int(size_text)
        if extents[0][0] != lba:
            raise ProbeValidationError(f"{context}: first extent does not match LBA")
        if sum(extent_size for _extent_lba, extent_size in extents) != size:
            raise ProbeValidationError(f"{context}: extents do not sum to size")
        entries.append(ManifestEntry(path=path, lba=lba, size=size, sha1=sha1))
    return tuple(entries)


def _normalize_extension(extension: str) -> str:
    normalized = extension.strip()
    if normalized.startswith("."):
        normalized = normalized[1:]
    if not normalized or re.fullmatch(r"[A-Za-z0-9]+", normalized) is None:
        raise ProbeValidationError("extension must contain only ASCII letters and digits")
    return normalized.upper()


def _entry_extension(path: str) -> str | None:
    leaf = ISO_VERSION_SUFFIX.sub("", path.rsplit("/", 1)[-1])
    if "." not in leaf:
        return None
    return leaf.rsplit(".", 1)[1].upper()


def _resolve_root(root_path: str | Path) -> Path:
    root = Path(root_path)
    if root.is_symlink():
        raise ProbeValidationError("extraction root must not be a symbolic link")
    try:
        resolved = root.resolve(strict=True)
    except OSError as exc:
        raise ProbeValidationError(f"cannot resolve extraction root: {exc}") from exc
    if not resolved.is_dir():
        raise ProbeValidationError("extraction root is not a directory")
    return resolved


def _resolve_candidate(root: Path, entry: ManifestEntry) -> Path:
    candidate = root / entry.path
    if candidate.is_symlink():
        raise ProbeValidationError("candidate must not be a symbolic link")
    try:
        resolved = candidate.resolve(strict=True)
    except OSError as exc:
        raise ProbeValidationError(f"cannot resolve candidate: {exc}") from exc
    try:
        resolved.relative_to(root)
    except ValueError as exc:
        raise ProbeValidationError("candidate escapes extraction root") from exc
    if not resolved.is_file():
        raise ProbeValidationError("candidate is not a regular file")
    try:
        actual_size = resolved.stat().st_size
    except OSError as exc:
        raise ProbeValidationError(f"cannot stat candidate: {exc}") from exc
    if actual_size != entry.size:
        raise ProbeValidationError("candidate size does not match manifest")
    return resolved


def _read_verified_prefix(candidate: Path, entry: ManifestEntry, prefix_size: int) -> bytes:
    """Read a bounded prefix while streaming a complete manifest-integrity check."""
    digest = hashlib.sha1()
    prefix = bytearray()
    try:
        with candidate.open("rb") as source:
            while chunk := source.read(64 * 1024):
                digest.update(chunk)
                remaining = prefix_size - len(prefix)
                if remaining > 0:
                    prefix.extend(chunk[:remaining])
    except OSError as exc:
        raise ProbeValidationError(f"cannot read candidate: {exc}") from exc
    if digest.hexdigest() != entry.sha1:
        raise ProbeValidationError("candidate SHA-1 does not match manifest")
    return bytes(prefix)


def _entropy_milli(data: bytes) -> int:
    if not data:
        return 0
    counts = Counter(data)
    entropy = -sum(
        (count / len(data)) * math.log2(count / len(data)) for count in counts.values()
    )
    return round(entropy * 1000)


def _monotonic_u32_run_count(prefix: bytes, file_size: int) -> int:
    """Count an intentionally generic offset-table heuristic in one prefix."""
    count = 0
    for offset in range(0, len(prefix) - 15, 4):
        words = tuple(
            int.from_bytes(prefix[offset + index : offset + index + 4], "little")
            for index in range(0, 16, 4)
        )
        if all(word < file_size for word in words) and words == tuple(sorted(words)):
            count += 1
    return count


def probe_prefix_features(
    manifest_path: str | Path,
    root_path: str | Path,
    *,
    extension: str,
    prefix_size: int = 4096,
) -> PrefixFeatureReport:
    """Measure opaque fixed-prefix features for one manifest-selected extension."""
    if not 1 <= prefix_size <= MAX_PREFIX_SIZE:
        raise ProbeValidationError(
            f"prefix size must be between 1 and {MAX_PREFIX_SIZE} bytes"
        )
    normalized_extension = _normalize_extension(extension)
    root = _resolve_root(root_path)
    entries = sorted(
        (entry for entry in read_manifest(manifest_path) if _entry_extension(entry.path) == normalized_extension),
        key=lambda entry: entry.path,
    )
    if not entries:
        raise ProbeValidationError("manifest has no candidates for the requested extension")

    prefixes: list[bytes] = []
    candidates: list[CandidateFeatures] = []
    for index, entry in enumerate(entries, start=1):
        candidate = _resolve_candidate(root, entry)
        prefix = _read_verified_prefix(candidate, entry, prefix_size)
        prefixes.append(prefix)
        candidates.append(
            CandidateFeatures(
                identifier=f"C{index:03d}",
                size=entry.size,
                prefix_size=len(prefix),
                prefix_sha1=hashlib.sha1(prefix).hexdigest(),
                entropy_milli=_entropy_milli(prefix),
                zero_bytes=prefix.count(0),
                monotonic_u32_run_count=_monotonic_u32_run_count(prefix, entry.size),
            )
        )

    common_prefix_span = min(len(prefix) for prefix in prefixes)
    constant_positions = sum(
        all(prefix[position] == prefixes[0][position] for prefix in prefixes)
        for position in range(common_prefix_span)
    )
    return PrefixFeatureReport(
        extension=normalized_extension,
        candidate_count=len(candidates),
        common_prefix_span=common_prefix_span,
        constant_positions=constant_positions,
        candidates=tuple(candidates),
    )


def format_prefix_features(report: PrefixFeatureReport) -> str:
    """Format a deterministic report without candidate paths or raw bytes."""
    lines = [
        "format\tsf3_probe-prefix-features-v1",
        f"candidate_count\t{report.candidate_count}",
        f"common_prefix_span\t{report.common_prefix_span}",
        f"constant_positions_0_common_prefix_span\t{report.constant_positions}",
    ]
    for candidate in report.candidates:
        lines.extend(
            [
                f"candidate\t{candidate.identifier}\tsize\t{candidate.size}",
                f"candidate\t{candidate.identifier}\tprefix_size\t{candidate.prefix_size}",
                f"candidate\t{candidate.identifier}\tprefix_sha1\t{candidate.prefix_sha1}",
                f"candidate\t{candidate.identifier}\tentropy_milli\t{candidate.entropy_milli}",
                f"candidate\t{candidate.identifier}\tzero_bytes\t{candidate.zero_bytes}",
                "candidate\t"
                f"{candidate.identifier}\tmonotonic_u32_run_count\t"
                f"{candidate.monotonic_u32_run_count}",
            ]
        )
    return "\n".join(lines)


def _positive_prefix_size(value: str) -> int:
    try:
        prefix_size = int(value, 10)
    except ValueError as exc:
        raise argparse.ArgumentTypeError("must be an integer") from exc
    if not 1 <= prefix_size <= MAX_PREFIX_SIZE:
        raise argparse.ArgumentTypeError(
            f"must be between 1 and {MAX_PREFIX_SIZE}"
        )
    return prefix_size


def main(argv: Sequence[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description="Produce opaque, read-only structural feature reports for manifest-tracked files."
    )
    commands = parser.add_subparsers(dest="command", required=True)
    prefix_parser = commands.add_parser(
        "prefix-features",
        help="measure fixed-prefix heuristics for one manifest-selected extension",
    )
    prefix_parser.add_argument("manifest", type=Path, help="validated extraction manifest")
    prefix_parser.add_argument("root", type=Path, help="matching extraction root")
    prefix_parser.add_argument(
        "--extension",
        required=True,
        help="extension class to select, without an ISO version suffix",
    )
    prefix_parser.add_argument(
        "--prefix-size",
        type=_positive_prefix_size,
        default=4096,
        help="bytes read from each candidate (default: 4096)",
    )
    args = parser.parse_args(argv)
    try:
        report = probe_prefix_features(
            args.manifest,
            args.root,
            extension=args.extension,
            prefix_size=args.prefix_size,
        )
    except ProbeValidationError as exc:
        print(f"sf3_probe: error: {exc}", file=sys.stderr)
        return 2
    print(format_prefix_features(report))
    return 0


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