#!/usr/bin/env python3
"""
HLDS/LG Stage 2 normalizer v3.

Stage 2 consumes Stage 1 output and performs conservative normalization only:

    * validate DLD containers
    * split valid DLD containers into firmware + optional stub
    * parse copied PRM files
    * match PRM entries to staged BIN payloads by package/filename/size
    * copy raw BINs into normalized_payloads
    * emit stage2_payload_manifest.json + stage2_inventory.csv

Stage 2 intentionally does NOT:
    * decrypt
    * perform MN103S opcode-density analysis
    * guess plaintext/encrypted status from byte-level heuristics
    * validate footer/checksum algorithms
    * generate final fleet profiles
"""

from __future__ import annotations

import argparse
import csv
import datetime as _dt
import hashlib
import json
import logging
import os
import re
import shutil
import struct
import sys
from collections import defaultdict
from dataclasses import dataclass, asdict
from pathlib import Path, PureWindowsPath
from typing import Any, Dict, Iterable, List, Optional, Tuple


PRODUCER = "hlds_stage2_normalizer_v3"
DLD_MAGICS = {b"HL-DT-ST", b"HL-DL-ST"}


def utc_iso() -> str:
    return _dt.datetime.now(_dt.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


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


def sha256_file(path: Path) -> str:
    h = hashlib.sha256()
    with path.open("rb") as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b""):
            h.update(chunk)
    return h.hexdigest()


def sanitize(s: str, fallback: str = "UNK") -> str:
    s = (s or "").strip()
    s = re.sub(r"[^A-Za-z0-9._-]+", "_", s)
    s = re.sub(r"_+", "_", s).strip("._-")
    return s or fallback


def rel_to_path(workspace: Path, rel_or_abs: str) -> Path:
    """Resolve Stage 1 JSON paths, which are usually Windows-style relative paths."""
    if not rel_or_abs:
        return workspace
    # Native absolute path first.
    p_native = Path(rel_or_abs)
    if p_native.is_absolute():
        return p_native

    # Windows absolute path on Windows; on POSIX this remains non-native but harmless
    # for script authoring. The script is intended for the same Windows workspace as Stage 1.
    p_win = PureWindowsPath(rel_or_abs)
    if p_win.is_absolute():
        return Path(str(p_win))

    return workspace.joinpath(*p_win.parts)


def relpath_str(path: Path, workspace: Path) -> str:
    try:
        rel = path.resolve().relative_to(workspace.resolve())
    except Exception:
        try:
            rel = Path(os.path.relpath(path, workspace))
        except Exception:
            rel = path
    return str(rel).replace("/", "\\")


def basename_rel(rel_or_path: str) -> str:
    return PureWindowsPath(rel_or_path).name


def parent_rel(rel_or_path: str) -> str:
    p = PureWindowsPath(rel_or_path)
    parent = str(p.parent)
    return "" if parent == "." else parent


def be32(b: bytes) -> int:
    return struct.unpack(">I", b)[0]


def le32(b: bytes) -> int:
    return struct.unpack("<I", b)[0]


def role_from_prm_filename(filename: str) -> Tuple[str, List[str]]:
    """Return firmware/stub/unknown from common HLDS PRM BIN naming patterns.

    Examples:
        DC10600M.BIN   -> firmware
        DC10600C.BIN   -> stub
        1SMA203C.BIN   -> firmware
        1SCA203C.BIN   -> stub
        16MA110C.BIN   -> firmware
        16CA110C.BIN   -> stub
    """
    stem = Path(filename).stem.upper()
    evidence: List[str] = []

    # Numeric-prefix forms where the M/C role letter is not the final character.
    # Observed examples: 1SMA203C / 1SCA203C, 16MA110C / 16CA110C.
    if re.match(r"^\d+S?M", stem):
        evidence.append("prm_filename_role_pattern_numeric_m")
        return "firmware", evidence
    if re.match(r"^\d+S?C", stem):
        evidence.append("prm_filename_role_pattern_numeric_c")
        return "stub", evidence

    # Simple suffix forms.
    if stem.endswith("M"):
        evidence.append("prm_filename_role_suffix_m")
        return "firmware", evidence
    if stem.endswith("C"):
        evidence.append("prm_filename_role_suffix_c")
        return "stub", evidence

    evidence.append("prm_filename_role_unknown")
    return "unknown", evidence


@dataclass
class PrmEntry:
    prm_id: str
    prm_copied_rel: str
    prm_model_hint: str
    source_root_file_rels: List[str]
    source_path_rels: List[str]
    line_no: int
    filename: str
    filename_upper: str
    expected_size: int
    expected_size_hex: str
    role_hint: str
    role_evidence: List[str]


@dataclass
class PrmFile:
    prm_id: str
    copied_rel: str
    sha256: str
    size: int
    model_hint: str
    sources: List[Dict[str, Any]]
    entries: List[Dict[str, Any]]
    parse_errors: List[str]


class Stage2:
    def __init__(self, args: argparse.Namespace):
        self.args = args
        self.workspace = Path(args.workspace).resolve()
        self.stage1_index = Path(args.stage1_index) if args.stage1_index else self.workspace / "stage1_payload_index.json"
        if not self.stage1_index.is_absolute():
            self.stage1_index = self.workspace / self.stage1_index

        self.normalized_dir = self.workspace / "normalized_payloads"
        self.logs_dir = self.workspace / "logs"
        self.log_path: Optional[Path] = None
        self.log = logging.getLogger(PRODUCER)

        self.counts: Dict[str, int] = {
            "stage1_payloads_seen": 0,
            "stage1_prm_files_seen": 0,
            "prm_files_parsed": 0,
            "prm_entries_parsed": 0,
            "prm_parse_errors": 0,
            "prm_bin_matches": 0,
            "prm_match_records": 0,
            "prm_matches_total": 0,
            "prm_role_conflicts": 0,
            "dld_seen": 0,
            "dld_valid": 0,
            "dld_invalid": 0,
            "dld_firmware_outputs": 0,
            "dld_stub_outputs": 0,
            "raw_bin_seen": 0,
            "raw_bin_outputs": 0,
            "missing_input_files": 0,
            "errors": 0,
        }

        self.stage1: Dict[str, Any] = {}
        self.prm_files: List[PrmFile] = []
        self.prm_entries: List[PrmEntry] = []
        self.records: List[Dict[str, Any]] = []
        self.inventory_rows: List[Dict[str, Any]] = []

    def setup(self) -> None:
        if self.args.clean:
            if self.normalized_dir.exists():
                shutil.rmtree(self.normalized_dir)
        self.normalized_dir.mkdir(parents=True, exist_ok=True)
        self.logs_dir.mkdir(parents=True, exist_ok=True)

        ts = _dt.datetime.now().strftime("%Y%m%d_%H%M%S")
        self.log_path = self.logs_dir / f"hlds_stage2_normalizer_{ts}.log"

        self.log.setLevel(logging.DEBUG)
        self.log.handlers.clear()

        file_handler = logging.FileHandler(self.log_path, encoding="utf-8")
        file_handler.setLevel(logging.DEBUG if self.args.verbose else logging.INFO)
        file_handler.setFormatter(logging.Formatter("[%(asctime)s] [%(levelname)s] %(message)s", "%H:%M:%S"))

        console_handler = logging.StreamHandler()
        console_handler.setLevel(logging.DEBUG if self.args.verbose else logging.INFO)
        console_handler.setFormatter(logging.Formatter("[%(levelname)s] %(message)s"))

        self.log.addHandler(file_handler)
        self.log.addHandler(console_handler)

    def load_stage1(self) -> None:
        self.log.info(f"loading Stage 1 index: {self.stage1_index}")
        with self.stage1_index.open("r", encoding="utf-8") as f:
            self.stage1 = json.load(f)
        self.counts["stage1_payloads_seen"] = len(self.stage1.get("payloads", []))
        self.counts["stage1_prm_files_seen"] = len(self.stage1.get("side_artifacts", {}).get("prm", []))

    def run(self) -> None:
        self.setup()
        self.load_stage1()
        self.parse_prms()
        self.process_payloads()
        self.write_outputs()
        self.print_summary()

    # -------------------------
    # PRM parsing / matching
    # -------------------------

    def parse_prms(self) -> None:
        for prm_item in self.stage1.get("side_artifacts", {}).get("prm", []):
            copied_rel = prm_item.get("copied_rel", "")
            prm_path = rel_to_path(self.workspace, copied_rel)
            prm_id = prm_item.get("sha256", "")[:16] or sanitize(Path(copied_rel).stem)
            parse_errors: List[str] = []
            entries: List[Dict[str, Any]] = []

            if not prm_path.exists():
                self.counts["missing_input_files"] += 1
                self.log.warning(f"PRM file missing: {copied_rel}")
                continue

            try:
                text = prm_path.read_text(encoding="ascii", errors="replace")
                lines = text.splitlines()
                parsed_entries: List[PrmEntry] = []

                source_root_file_rels = sorted({
                    s.get("root_file_rel", "")
                    for s in prm_item.get("sources", [])
                    if s.get("root_file_rel")
                })
                source_path_rels = [
                    s.get("source_path_rel", "")
                    for s in prm_item.get("sources", [])
                    if s.get("source_path_rel")
                ]

                for line_no, line in enumerate(lines, 1):
                    raw = line.strip()
                    if not raw:
                        continue
                    parts = re.split(r"\s+", raw)
                    if len(parts) < 2:
                        parse_errors.append(f"line {line_no}: expected '<filename> <size>', got {raw!r}")
                        continue

                    filename = parts[0].strip()
                    try:
                        expected_size = int(parts[1], 0)
                    except ValueError:
                        parse_errors.append(f"line {line_no}: invalid size {parts[1]!r}")
                        continue

                    role_hint, role_evidence = role_from_prm_filename(filename)
                    entry = PrmEntry(
                        prm_id=prm_id,
                        prm_copied_rel=copied_rel,
                        prm_model_hint=prm_item.get("model_hint", ""),
                        source_root_file_rels=source_root_file_rels,
                        source_path_rels=source_path_rels,
                        line_no=line_no,
                        filename=filename,
                        filename_upper=filename.upper(),
                        expected_size=expected_size,
                        expected_size_hex=f"0x{expected_size:X}",
                        role_hint=role_hint,
                        role_evidence=role_evidence,
                    )
                    parsed_entries.append(entry)

                # If role is still unknown and the PRM has exactly two distinct sizes,
                # use size ranking as a weak hint: larger=firmware, smaller=stub.
                if len(parsed_entries) == 2:
                    e0, e1 = parsed_entries
                    if e0.expected_size != e1.expected_size:
                        large, small = (e0, e1) if e0.expected_size > e1.expected_size else (e1, e0)
                        if large.role_hint == "unknown":
                            large.role_hint = "firmware"
                            large.role_evidence.append("prm_pair_larger_size_weak")
                        if small.role_hint == "unknown":
                            small.role_hint = "stub"
                            small.role_evidence.append("prm_pair_smaller_size_weak")

                for e in parsed_entries:
                    self.prm_entries.append(e)
                    entries.append(asdict(e))
                    self.counts["prm_entries_parsed"] += 1

                if parse_errors:
                    self.counts["prm_parse_errors"] += len(parse_errors)
                    for err in parse_errors:
                        self.log.warning(f"PRM parse issue {copied_rel}: {err}")

                self.counts["prm_files_parsed"] += 1
                self.log.info(f"parsed PRM: {copied_rel} entries={len(parsed_entries)}")

                self.prm_files.append(PrmFile(
                    prm_id=prm_id,
                    copied_rel=copied_rel,
                    sha256=prm_item.get("sha256", ""),
                    size=prm_item.get("size", prm_path.stat().st_size),
                    model_hint=prm_item.get("model_hint", ""),
                    sources=prm_item.get("sources", []),
                    entries=entries,
                    parse_errors=parse_errors,
                ))
            except Exception as e:
                self.counts["errors"] += 1
                self.log.error(f"failed to parse PRM {copied_rel}: {type(e).__name__}: {e}")

    def find_prm_matches_for_bin(self, payload: Dict[str, Any]) -> List[Dict[str, Any]]:
        """Return every PRM entry that matches this BIN.

        Stage 1 dedupes by SHA-256, so one staged BIN can represent multiple
        source packages and therefore multiple PRM entries. Stage 2 must retain
        all of those PRM matches rather than choosing the first one.
        """
        if payload.get("ext", "").lower() != ".bin" or not self.prm_entries:
            return []

        size = int(payload.get("size", -1))
        sources = payload.get("sources", []) or []
        source_names = {basename_rel(s.get("source_path_rel", "")).upper() for s in sources}
        source_roots = {s.get("root_file_rel", "") for s in sources if s.get("root_file_rel")}

        def pack(entry: PrmEntry, strength: str, matched_by: str, source_root: str = "") -> Dict[str, Any]:
            return {
                "matched": True,
                "match_strength": strength,
                "matched_by": matched_by,
                "matched_source_root": source_root,
                "role_hint": entry.role_hint,
                "expected_size": entry.expected_size,
                "expected_size_hex": entry.expected_size_hex,
                "prm_filename": entry.filename,
                "prm_copied_rel": entry.prm_copied_rel,
                "prm_id": entry.prm_id,
                "role_evidence": entry.role_evidence,
            }

        matches: List[Dict[str, Any]] = []
        seen: set[Tuple[str, str, int, str]] = set()

        def add_match(m: Dict[str, Any]) -> None:
            # Dedupe by the actual PRM entry, not by how it was matched.
            # v2 could record the same PRM line twice, once as a strong same-package
            # match and once again as a weak global match. That inflated match counts.
            key = (
                m.get("prm_id", ""),
                m.get("prm_filename", "").upper(),
                int(m.get("expected_size", -1)),
            )
            if key not in seen:
                seen.add(key)
                matches.append(m)

        # Strong: same package/root + exact source filename + exact size.
        for root in sorted(source_roots):
            for entry in self.prm_entries:
                if (
                    root in entry.source_root_file_rels
                    and entry.expected_size == size
                    and entry.filename_upper in source_names
                ):
                    add_match(pack(entry, "strong", "same_package_filename_size", root))

        # Good: same package/root + exact size, and unique among PRM entries for that root.
        # Only add this if a strong match was not already found for that root.
        roots_with_strong = {
            m.get("matched_source_root", "")
            for m in matches
            if m.get("match_strength") == "strong"
        }
        for root in sorted(source_roots):
            if root in roots_with_strong:
                continue
            candidates = [
                e for e in self.prm_entries
                if root in e.source_root_file_rels and e.expected_size == size
            ]
            if len(candidates) == 1:
                add_match(pack(candidates[0], "good", "same_package_size_unique", root))

        # Weak but still useful: global exact filename + exact size, unique across PRM entries.
        # Only use this for source filenames not already covered by same-package evidence.
        candidates = [
            e for e in self.prm_entries
            if e.expected_size == size and e.filename_upper in source_names
        ]
        unique = {(e.prm_id, e.filename_upper, e.expected_size) for e in candidates}
        if len(unique) == 1 and candidates:
            add_match(pack(candidates[0], "weak", "global_filename_size_unique", ""))

        # Sort for deterministic output: strong, good, weak, then PRM path/name.
        strength_order = {"strong": 0, "good": 1, "weak": 2}
        matches.sort(key=lambda m: (
            strength_order.get(m.get("match_strength", ""), 99),
            m.get("prm_copied_rel", ""),
            m.get("prm_filename", ""),
        ))
        return matches

    def summarize_prm_matches(self, matches: List[Dict[str, Any]]) -> Tuple[str, Dict[str, Any], List[str]]:
        """Derive a conservative role and summary from all PRM matches."""
        if not matches:
            return "unknown", {
                "matched": False,
                "match_count": 0,
                "role_agreement": "none",
            }, []

        roles = [m.get("role_hint", "unknown") for m in matches]
        concrete_roles = [r for r in roles if r in {"firmware", "stub"}]
        notes: List[str] = []

        if len(concrete_roles) == len(roles) and len(set(concrete_roles)) == 1:
            role = concrete_roles[0]
            role_agreement = "all_agree"
        elif len(set(concrete_roles)) > 1:
            role = "unknown"
            role_agreement = "conflict"
            notes.append("prm_role_conflict")
        else:
            role = "unknown"
            role_agreement = "inconclusive"
            notes.append("prm_role_inconclusive")

        strengths = sorted({m.get("match_strength", "") for m in matches if m.get("match_strength")})
        summary = {
            "matched": True,
            "match_count": len(matches),
            "role_agreement": role_agreement,
            "selected_role_hint": role,
            "roles_seen": sorted(set(roles)),
            "match_strengths": strengths,
        }
        return role, summary, notes

    # -------------------------
    # DLD parsing / splitting
    # -------------------------

    def parse_dld_header(self, data: bytes) -> Dict[str, Any]:
        if len(data) < 0x40:
            return {"valid": False, "errors": ["too_short_for_dld_header"]}

        magic = data[0x08:0x10]
        errors: List[str] = []
        warnings: List[str] = []

        if magic not in DLD_MAGICS:
            errors.append(f"bad_magic:{magic!r}")

        format_version = le32(data[0x04:0x08])
        if format_version != 1:
            warnings.append(f"unexpected_format_version:{format_version}")

        firmware_block_size = be32(data[0x10:0x14])
        payload_size = be32(data[0x1C:0x20])
        firmware_size = firmware_block_size - 0x40
        stub_size = payload_size - firmware_size
        dld_size = payload_size + 0x40

        if firmware_block_size < 0x40:
            errors.append("firmware_block_size_less_than_header")
        if firmware_size <= 0:
            errors.append("nonpositive_firmware_size")
        if payload_size < firmware_size:
            errors.append("payload_size_less_than_firmware_size")
        if stub_size < 0:
            errors.append("negative_stub_size")
        if dld_size > len(data):
            errors.append("declared_dld_size_exceeds_file_size")
        if dld_size < len(data):
            warnings.append(f"trailing_bytes_after_declared_dld_size:{len(data) - dld_size}")

        identity_raw = data[0x20:0x40]
        identity = identity_raw.split(b"\x00", 1)[0].decode("ascii", errors="replace").strip()

        return {
            "valid": not errors,
            "errors": errors,
            "warnings": warnings,
            "magic": magic.decode("ascii", errors="replace"),
            "format_version": format_version,
            "firmware_block_size": firmware_block_size,
            "firmware_size": firmware_size,
            "payload_size": payload_size,
            "stub_size": stub_size,
            "dld_size": dld_size,
            "identity": identity,
            "offsets": {
                "header": {"offset": 0, "size": 0x40},
                "firmware": {"offset": 0x40, "size": firmware_size},
                "stub": {"offset": 0x40 + firmware_size, "size": max(stub_size, 0)},
            },
        }

    # -------------------------
    # Payload processing
    # -------------------------

    def process_payloads(self) -> None:
        for payload in self.stage1.get("payloads", []):
            ext = payload.get("ext", "").lower()
            if ext == ".dld":
                self.process_dld(payload)
            elif ext == ".bin":
                self.process_bin(payload)
            else:
                self.log.debug(f"Stage 1 payload with unsupported ext ignored: {payload.get('staged_rel')}")

    def payload_path(self, payload: Dict[str, Any]) -> Path:
        return rel_to_path(self.workspace, payload.get("staged_rel", ""))

    def base_for_payload(self, payload: Dict[str, Any]) -> str:
        identity = payload.get("identity_hint", {})
        model = sanitize(identity.get("model", "UNK"))
        version = sanitize(identity.get("version", "UNK"))
        pid = sanitize(payload.get("payload_id", payload.get("sha256", "")[:16]))[:8]
        return f"{model}_{version}_{pid}"

    def write_component(self, data: bytes, base: str, suffix: str, ext: str = ".bin") -> Tuple[Path, str]:
        digest = sha256_bytes(data)
        dest = self.normalized_dir / f"{base}_{suffix}_{digest[:8]}{ext}"
        if not dest.exists():
            dest.write_bytes(data)
        return dest, digest

    def common_source_info(self, payload: Dict[str, Any]) -> Dict[str, Any]:
        return {
            "stage1_payload_id": payload.get("payload_id"),
            "stage1_staged_rel": payload.get("staged_rel"),
            "stage1_ext": payload.get("ext"),
            "stage1_sha256": payload.get("sha256"),
            "stage1_first_source": payload.get("first_source", {}),
            "stage1_sources": payload.get("sources", []),
            "identity_hint": payload.get("identity_hint", {}),
        }

    def add_record(self, record: Dict[str, Any]) -> None:
        self.records.append(record)
        self.inventory_rows.append({
            "record_id": record.get("record_id", ""),
            "record_type": record.get("record_type", ""),
            "role": record.get("role", ""),
            "normalized_rel": record.get("normalized_rel", ""),
            "size": record.get("size", ""),
            "sha256": record.get("sha256", ""),
            "source_stage1_payload_id": record.get("source", {}).get("stage1_payload_id", ""),
            "source_staged_rel": record.get("source", {}).get("stage1_staged_rel", ""),
            "model": record.get("source", {}).get("identity_hint", {}).get("model", ""),
            "version": record.get("source", {}).get("identity_hint", {}).get("version", ""),
            "evidence": ";".join(record.get("evidence", [])),
            "notes": ";".join(record.get("notes", [])),
        })

    def process_dld(self, payload: Dict[str, Any]) -> None:
        self.counts["dld_seen"] += 1
        path = self.payload_path(payload)
        source = self.common_source_info(payload)
        base = self.base_for_payload(payload)

        if not path.exists():
            self.counts["missing_input_files"] += 1
            self.log.warning(f"DLD staged file missing: {payload.get('staged_rel')}")
            return

        try:
            data = path.read_bytes()
            info = self.parse_dld_header(data)

            container_record = {
                "record_id": f"dld_{payload.get('payload_id', sha256_bytes(data)[:16])}",
                "record_type": "container",
                "role": "container",
                "container_type": "dld",
                "normalized_rel": "",
                "size": len(data),
                "sha256": sha256_bytes(data),
                "source": source,
                "dld": info,
                "components": [],
                "evidence": ["stage1_dld_payload", "dld_header_parsed"],
                "notes": [],
                "analysis_status": "pending_stage3",
            }

            if not info.get("valid"):
                self.counts["dld_invalid"] += 1
                container_record["evidence"].append("dld_header_invalid")
                container_record["notes"].extend(info.get("errors", []))
                self.add_record(container_record)
                self.log.warning(f"invalid DLD: {payload.get('staged_rel')} errors={info.get('errors')}")
                return

            self.counts["dld_valid"] += 1
            if info.get("warnings"):
                container_record["notes"].extend(info["warnings"])

            dld_size = info["dld_size"]
            firmware_size = info["firmware_size"]
            stub_size = info["stub_size"]

            header = data[0:0x40]
            firmware = data[0x40:0x40 + firmware_size]
            stub = data[0x40 + firmware_size:0x40 + firmware_size + stub_size] if stub_size > 0 else b""

            header_path, header_sha = self.write_component(header, base, "dld_header")
            container_record["components"].append({
                "role": "container",
                "kind": "dld_header",
                "normalized_rel": relpath_str(header_path, self.workspace),
                "size": len(header),
                "sha256": header_sha,
            })

            fw_path, fw_sha = self.write_component(firmware, base, "firmware")
            fw_record = {
                "record_id": fw_sha[:16],
                "record_type": "normalized_payload",
                "role": "firmware",
                "container_type": "dld_component",
                "normalized_rel": relpath_str(fw_path, self.workspace),
                "size": len(firmware),
                "sha256": fw_sha,
                "source": source,
                "dld_component": {
                    "parent_dld_sha256": sha256_bytes(data),
                    "parent_dld_declared_size": dld_size,
                    "offset": 0x40,
                    "size": len(firmware),
                },
                "prm_match_summary": {"matched": False, "match_count": 0, "role_agreement": "none"},
                "prm_matches": [],
                "evidence": ["dld_header_firmware_size", "dld_split_layout"],
                "notes": [],
                "analysis_status": "pending_stage3",
            }
            self.add_record(fw_record)
            self.counts["dld_firmware_outputs"] += 1
            container_record["components"].append({
                "role": "firmware",
                "kind": "dld_firmware",
                "normalized_rel": fw_record["normalized_rel"],
                "size": fw_record["size"],
                "sha256": fw_sha,
            })

            if stub_size > 0:
                stub_path, stub_sha = self.write_component(stub, base, "stub")
                stub_record = {
                    "record_id": stub_sha[:16],
                    "record_type": "normalized_payload",
                    "role": "stub",
                    "container_type": "dld_component",
                    "normalized_rel": relpath_str(stub_path, self.workspace),
                    "size": len(stub),
                    "sha256": stub_sha,
                    "source": source,
                    "dld_component": {
                        "parent_dld_sha256": sha256_bytes(data),
                        "parent_dld_declared_size": dld_size,
                        "offset": 0x40 + firmware_size,
                        "size": len(stub),
                    },
                    "prm_match_summary": {"matched": False, "match_count": 0, "role_agreement": "none"},
                    "prm_matches": [],
                    "evidence": ["dld_header_stub_size", "dld_split_layout"],
                    "notes": [],
                    "analysis_status": "pending_stage3",
                }
                self.add_record(stub_record)
                self.counts["dld_stub_outputs"] += 1
                container_record["components"].append({
                    "role": "stub",
                    "kind": "dld_stub",
                    "normalized_rel": stub_record["normalized_rel"],
                    "size": stub_record["size"],
                    "sha256": stub_sha,
                })
            else:
                container_record["evidence"].append("dld_no_stub_declared")

            self.add_record(container_record)
            self.log.info(
                f"DLD normalized: {payload.get('staged_rel')} firmware=0x{len(firmware):X} stub=0x{len(stub):X}"
            )
        except Exception as e:
            self.counts["errors"] += 1
            self.log.error(f"failed DLD processing {payload.get('staged_rel')}: {type(e).__name__}: {e}")

    def process_bin(self, payload: Dict[str, Any]) -> None:
        self.counts["raw_bin_seen"] += 1
        path = self.payload_path(payload)
        source = self.common_source_info(payload)
        base = self.base_for_payload(payload)

        if not path.exists():
            self.counts["missing_input_files"] += 1
            self.log.warning(f"BIN staged file missing: {payload.get('staged_rel')}")
            return

        try:
            prm_matches = self.find_prm_matches_for_bin(payload)
            role, prm_match_summary, prm_notes = self.summarize_prm_matches(prm_matches)

            data = path.read_bytes()
            digest = sha256_bytes(data)
            dest = self.normalized_dir / f"{base}_{role}_{digest[:8]}.bin"
            if not dest.exists():
                shutil.copy2(path, dest)

            evidence = ["stage1_raw_bin_payload"]
            if prm_matches:
                for strength in sorted({m.get("match_strength", "") for m in prm_matches if m.get("match_strength")}):
                    evidence.append(f"prm_match_{strength}")
                for matched_by in sorted({m.get("matched_by", "") for m in prm_matches if m.get("matched_by")}):
                    evidence.append(matched_by)
                evidence.append(f"prm_match_count_{len(prm_matches)}")

                self.counts["prm_bin_matches"] += 1
                self.counts["prm_match_records"] += 1
                self.counts["prm_matches_total"] += len(prm_matches)
                if prm_match_summary.get("role_agreement") == "conflict":
                    self.counts["prm_role_conflicts"] += 1
            else:
                evidence.append("no_prm_match")

            record = {
                "record_id": digest[:16],
                "record_type": "normalized_payload",
                "role": role,
                "container_type": "raw_bin",
                "normalized_rel": relpath_str(dest, self.workspace),
                "size": len(data),
                "sha256": digest,
                "source": source,
                "prm_match_summary": prm_match_summary,
                "prm_matches": prm_matches,
                "evidence": evidence,
                "notes": prm_notes,
                "analysis_status": "pending_stage3",
            }
            self.add_record(record)
            self.counts["raw_bin_outputs"] += 1
            self.log.info(
                f"BIN normalized: {payload.get('staged_rel')} role_hint={role} "
                f"size=0x{len(data):X} prm_matches={len(prm_matches)}"
            )
        except Exception as e:
            self.counts["errors"] += 1
            self.log.error(f"failed BIN processing {payload.get('staged_rel')}: {type(e).__name__}: {e}")

    # -------------------------
    # Output
    # -------------------------

    def write_outputs(self) -> None:
        manifest = {
            "schema_version": 1,
            "producer": PRODUCER,
            "created_utc": utc_iso(),
            "workspace": str(self.workspace),
            "stage1_index_rel": relpath_str(self.stage1_index, self.workspace),
            "normalized_dir": relpath_str(self.normalized_dir, self.workspace),
            "log_rel": relpath_str(self.log_path, self.workspace) if self.log_path else "",
            "counts": self.counts,
            "stage_boundaries": {
                "stage2_does": [
                    "validate_dld",
                    "split_dld",
                    "parse_prm",
                    "match_prm_to_bin_one_to_many",
                    "normalize_raw_bin",
                    "record_conservative_role_hints",
                ],
                "stage2_does_not": [
                    "decrypt",
                    "mn103s_opcode_density",
                    "stub_footer_checksum_analysis",
                    "plaintext_encrypted_guess",
                    "firmware_size_family_analysis",
                    "vector_reset_table_analysis",
                    "final_profile_generation",
                ],
            },
            "roles": ["firmware", "stub", "container", "unknown"],
            "prm_files": [asdict(p) for p in self.prm_files],
            "records": self.records,
        }

        manifest_path = self.workspace / "stage2_payload_manifest.json"
        manifest_path.write_text(json.dumps(manifest, indent=4), encoding="utf-8")
        self.log.info(f"wrote Stage 2 manifest: {relpath_str(manifest_path, self.workspace)}")

        inventory_path = self.workspace / "stage2_inventory.csv"
        fields = [
            "record_id",
            "record_type",
            "role",
            "normalized_rel",
            "size",
            "sha256",
            "source_stage1_payload_id",
            "source_staged_rel",
            "model",
            "version",
            "evidence",
            "notes",
        ]
        with inventory_path.open("w", encoding="utf-8", newline="") as f:
            w = csv.DictWriter(f, fieldnames=fields)
            w.writeheader()
            for row in self.inventory_rows:
                w.writerow(row)
        self.log.info(f"wrote Stage 2 inventory: {relpath_str(inventory_path, self.workspace)}")

    def print_summary(self) -> None:
        print()
        print("=" * 72)
        print("HLDS STAGE 2 NORMALIZER SUMMARY")
        print("=" * 72)
        print(f"Workspace          : {self.workspace}")
        print(f"Stage 1 index      : {self.stage1_index}")
        print(f"Normalized dir     : {self.normalized_dir}")
        print(f"Manifest           : {self.workspace / 'stage2_payload_manifest.json'}")
        print(f"Inventory CSV      : {self.workspace / 'stage2_inventory.csv'}")
        print(f"Log file           : {self.log_path}")
        print("-" * 72)
        print(f"Stage1 payloads    : {self.counts['stage1_payloads_seen']}")
        print(f"PRM files parsed   : {self.counts['prm_files_parsed']}")
        print(f"PRM entries parsed : {self.counts['prm_entries_parsed']}")
        print(f"PRM match records  : {self.counts['prm_match_records']}")
        print(f"PRM matches total  : {self.counts['prm_matches_total']}")
        print(f"PRM role conflicts : {self.counts['prm_role_conflicts']}")
        print(f"DLD seen           : {self.counts['dld_seen']}")
        print(f"DLD valid          : {self.counts['dld_valid']}")
        print(f"DLD invalid        : {self.counts['dld_invalid']}")
        print(f"DLD firmware outs  : {self.counts['dld_firmware_outputs']}")
        print(f"DLD stub outs      : {self.counts['dld_stub_outputs']}")
        print(f"Raw BIN outputs    : {self.counts['raw_bin_outputs']}")
        print(f"Missing files      : {self.counts['missing_input_files']}")
        print(f"Errors             : {self.counts['errors']}")
        print("=" * 72)


def build_arg_parser() -> argparse.ArgumentParser:
    ap = argparse.ArgumentParser(
        description="Stage 2 HLDS/LG normalizer: split DLDs, parse PRMs, match PRM entries to BINs, and write normalized payloads."
    )
    ap.add_argument(
        "--workspace",
        default=r"D:\vscode-workspace\xbox-dvd-tools-v2",
        help="Workspace root containing Stage 1 outputs.",
    )
    ap.add_argument(
        "--stage1-index",
        default=None,
        help="Path to stage1_payload_index.json. Defaults to <workspace>\\stage1_payload_index.json.",
    )
    ap.add_argument(
        "--clean",
        action="store_true",
        help="Delete normalized_payloads before writing new normalized outputs.",
    )
    ap.add_argument(
        "--verbose",
        action="store_true",
        help="Verbose console and file logging.",
    )
    return ap


def main(argv: Optional[List[str]] = None) -> int:
    args = build_arg_parser().parse_args(argv)
    s = Stage2(args)
    try:
        s.run()
        return 0 if s.counts.get("errors", 0) == 0 else 2
    except KeyboardInterrupt:
        print("Interrupted.", file=sys.stderr)
        return 130
    except Exception as e:
        print(f"fatal: {type(e).__name__}: {e}", file=sys.stderr)
        return 1


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