#!/usr/bin/env python3
"""
HLDS/LG Stage 5A MN103S disassembler.

This ports the earlier auto_disassembler.py idea into the staged pipeline.

Inputs:
    * stage5_analysis_inputs.json
    * plaintext firmware/stub binaries referenced by that manifest

v3 recognizes Stage 3/4 role provenance metadata:
    * target_scope: donor_xbox / validation_reference / recipient_candidate
    * original_role: original Stage 2/3 role before promotion
    * effective_role: role used by this stage for disassembly

Outputs:
    * disassembly/*.asm
    * stage5_disassembly_manifest.json
    * stage5_disassembly_inventory.csv
    * stage5_analyzer_inputs.json
    * logs/hlds_stage5_disassembler_YYYYMMDD_HHMMSS.log

This stage disassembles only. It does not search for command handlers,
checksum routines, flash routines, vector tables, or profiles yet.
Those belong to the Stage 5B analyzer that consumes this stage's output.

Default external disassembler command:
    wsl mn10300-elf-objdump -D -b binary -m mn10300 -EB --adjust-vma=<base> <filename>

Firmware VMA detection is ported from the previous script:
    scan 32-bit word lanes; compare count of high-byte 0x40 vs 0x90;
    choose 0x90000000 if 0x90 wins, otherwise 0x40000000.

Stub base address is fixed by default:
    0x40008000
"""

from __future__ import annotations

import argparse
import csv
import datetime as _dt
import hashlib
import json
import os
import re
import shutil
import subprocess
import sys
from pathlib import Path, PureWindowsPath
from typing import Any, Dict, List, Optional, Sequence, Tuple


PRODUCER = "hlds_stage5_disassembler_v3"


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 rel_to_path(workspace: Path, rel_or_abs: str) -> Path:
    if not rel_or_abs:
        return workspace

    p_native = Path(rel_or_abs)
    if p_native.is_absolute():
        return p_native

    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 sanitize_name(s: str) -> str:
    s = s.replace("\\", "_").replace("/", "_").replace(":", "_")
    s = re.sub(r"[^A-Za-z0-9._-]+", "_", s)
    s = re.sub(r"_+", "_", s)
    return s.strip("_") or "payload"


def detect_base_vma_programmatically(bin_path: Path) -> Dict[str, Any]:
    data = bin_path.read_bytes()
    count_40 = 0
    count_90 = 0

    # Ported intentionally from the earlier script:
    # for each 4-byte word, look at byte i+3 as the high byte under the way
    # pointer values appear in these binaries.
    for i in range(0, len(data) - 3, 4):
        msb = data[i + 3]
        if msb == 0x40:
            count_40 += 1
        elif msb == 0x90:
            count_90 += 1

    base = "0x90000000" if count_90 > count_40 else "0x40000000"
    confidence = "medium"
    total = count_40 + count_90
    if total == 0:
        confidence = "low"
    elif abs(count_40 - count_90) / total >= 0.50:
        confidence = "high"

    return {
        "base_addr": base,
        "count_40": count_40,
        "count_90": count_90,
        "confidence": confidence,
        "method": "word_high_byte_0x40_vs_0x90",
    }


def build_objdump_command(
    bin_filename: str,
    base_addr: str,
    objdump: str,
    use_wsl: bool,
    extra_args: Sequence[str],
) -> List[str]:
    cmd = []
    if use_wsl:
        cmd.append("wsl")
    cmd.extend([
        objdump,
        "-D",
        "-b", "binary",
        "-m", "mn10300",
        "-EB",
        f"--adjust-vma={base_addr}",
    ])
    cmd.extend(extra_args)
    cmd.append(bin_filename)
    return cmd


def target_scope_from_item(item: Dict[str, Any]) -> str:
    if item.get("target_scope"):
        return str(item.get("target_scope"))
    route = item.get("stage3_route", {}) or {}
    if route.get("target_scope"):
        return str(route.get("target_scope"))

    source = item.get("source", {}) or {}
    ident = source.get("identity_hint", {}) or {}
    text = " ".join([
        str(ident.get("model", "")),
        str(item.get("plaintext_rel", "")),
        str(item.get("payload_rel", "")),
    ]).upper().replace("-", "").replace("_", "")
    if "GDR8050L" in text:
        return "donor_xbox"
    if "GDR3120L" in text:
        return "validation_reference"
    return "recipient_candidate"



def original_role_from_item(item: Dict[str, Any]) -> str:
    """Preserve the first known role before Stage 4/5 promotion."""
    if item.get("original_role"):
        return str(item.get("original_role"))
    route = item.get("stage3_route", {}) or {}
    if route.get("original_role"):
        return str(route.get("original_role"))
    return str(item.get("role", "unknown"))


def effective_role_from_item(item: Dict[str, Any], role_class: Optional[str] = None) -> str:
    """Role to use for downstream analysis. This may differ from original_role."""
    if item.get("effective_role") in ("firmware", "stub", "unknown"):
        return str(item.get("effective_role"))
    if role_class in ("firmware", "stub", "unknown"):
        return str(role_class)
    role = item.get("role", "unknown")
    if role in ("firmware", "stub", "unknown"):
        return str(role)
    return "unknown"


def classify_input(item: Dict[str, Any]) -> str:
    role = item.get("role", "unknown")
    kind = item.get("input_kind", "")
    route = item.get("stage3_route", {}) or {}
    route_name = str(route.get("route", ""))

    if role == "stub" or "stub" in kind or "stub" in route_name:
        return "stub"
    if (
        role == "firmware"
        or "firmware" in kind
        or "firmware" in route_name
        or target_scope_from_item(item) == "validation_reference"
    ):
        return "firmware"
    return "unknown"


class Stage5Disassembler:
    def __init__(self, args: argparse.Namespace):
        self.args = args
        self.workspace = Path(args.workspace).resolve()

        self.inputs_path = Path(args.inputs) if args.inputs else self.workspace / "stage5_analysis_inputs.json"
        if not self.inputs_path.is_absolute():
            self.inputs_path = self.workspace / self.inputs_path

        self.output_dir = Path(args.output_dir) if args.output_dir else self.workspace / "disassembly"
        if not self.output_dir.is_absolute():
            self.output_dir = self.workspace / self.output_dir

        self.logs_dir = self.workspace / "logs"
        self.log_path: Optional[Path] = None

        self.inputs_doc: Dict[str, Any] = {}
        self.records: List[Dict[str, Any]] = []
        self.inventory_rows: List[Dict[str, Any]] = []
        self.analyzer_inputs: List[Dict[str, Any]] = []
        self.review_inputs: List[Dict[str, Any]] = []

        self.counts = {
            "analysis_inputs_seen": 0,
            "review_inputs_seen": 0,
            "inputs_selected": 0,
            "firmware_inputs": 0,
            "stub_inputs": 0,
            "unknown_inputs": 0,
            "disassembly_attempts": 0,
            "disassembly_succeeded": 0,
            "disassembly_failed": 0,
            "dry_run_records": 0,
            "missing_input_files": 0,
            "sha256_mismatches": 0,
            "original_role_unknown_promoted_to_firmware": 0,
            "errors": 0,
        }

    def log(self, msg: str) -> None:
        print(msg)
        if self.log_path:
            with self.log_path.open("a", encoding="utf-8") as f:
                f.write(msg + "\n")

    def setup(self) -> None:
        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_stage5_disassembler_{ts}.log"

        if self.args.clean and self.output_dir.exists():
            shutil.rmtree(self.output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)

        self.log(f"[INFO] {PRODUCER} started {utc_iso()}")
        if self.args.dry_run:
            self.log("[INFO] dry-run mode: commands will be recorded but not executed")

    def load_inputs(self) -> None:
        self.log(f"[INFO] loading Stage 5 analysis inputs: {self.inputs_path}")
        with self.inputs_path.open("r", encoding="utf-8") as f:
            self.inputs_doc = json.load(f)

        self.counts["analysis_inputs_seen"] = len(self.inputs_doc.get("analysis_inputs", []))
        self.counts["review_inputs_seen"] = len(self.inputs_doc.get("review_inputs", []))

    def run(self) -> None:
        self.setup()
        self.load_inputs()
        self.process_inputs()
        self.write_outputs()
        self.print_summary()

    def iter_selected_inputs(self) -> List[Dict[str, Any]]:
        selected = list(self.inputs_doc.get("analysis_inputs", []))
        if self.args.include_review:
            for item in self.inputs_doc.get("review_inputs", []):
                # review inputs use payload_rel/payload_sha256 rather than plaintext_rel/plaintext_sha256
                normalized = dict(item)
                if "plaintext_rel" not in normalized and "payload_rel" in normalized:
                    normalized["plaintext_rel"] = normalized.get("payload_rel")
                if "plaintext_sha256" not in normalized and "payload_sha256" in normalized:
                    normalized["plaintext_sha256"] = normalized.get("payload_sha256")
                normalized["input_kind"] = normalized.get("input_kind", "review")
                if "target_scope" not in normalized:
                    normalized["target_scope"] = target_scope_from_item(normalized)
                if "original_role" not in normalized:
                    normalized["original_role"] = original_role_from_item(normalized)
                # effective_role is finalized after classify_input(), but preserve if already provided.
                selected.append(normalized)

        # Optional role filters.
        allowed_roles = set(self.args.roles or [])
        if allowed_roles:
            selected = [x for x in selected if x.get("role") in allowed_roles]

        # Deduplicate exact file+hash pairs.
        seen = set()
        deduped = []
        for item in selected:
            key = (item.get("plaintext_rel"), item.get("plaintext_sha256"))
            if key in seen:
                continue
            seen.add(key)
            deduped.append(item)
        return deduped

    def process_inputs(self) -> None:
        selected = self.iter_selected_inputs()
        self.counts["inputs_selected"] = len(selected)

        for idx, item in enumerate(selected):
            self.process_one(idx, item)

    def process_one(self, idx: int, item: Dict[str, Any]) -> None:
        rel = item.get("plaintext_rel", "")
        expected_sha = item.get("plaintext_sha256", "")
        if "target_scope" not in item or "original_role" not in item:
            item = dict(item)
            item["target_scope"] = target_scope_from_item(item)
            item["original_role"] = original_role_from_item(item)
        role_class = classify_input(item)
        if "effective_role" not in item:
            item = dict(item)
            item["effective_role"] = effective_role_from_item(item, role_class)

        if role_class == "firmware":
            self.counts["firmware_inputs"] += 1
            if original_role_from_item(item) == "unknown":
                self.counts["original_role_unknown_promoted_to_firmware"] += 1
        elif role_class == "stub":
            self.counts["stub_inputs"] += 1
        else:
            self.counts["unknown_inputs"] += 1

        bin_path = rel_to_path(self.workspace, rel)
        if not bin_path.exists():
            self.counts["missing_input_files"] += 1
            rec = self.record_base(item, idx, role_class)
            rec.update({
                "status": "MISSING_INPUT",
                "message": f"input not found: {rel}",
            })
            self.records.append(rec)
            self.log(f"[ERROR] missing input: {rel}")
            return

        try:
            data = bin_path.read_bytes()
            actual_sha = sha256_bytes(data)
            notes: List[str] = []
            if expected_sha and expected_sha != actual_sha:
                self.counts["sha256_mismatches"] += 1
                notes.append("sha256_mismatch")

            base_info = self.choose_base_address(item, role_class, bin_path)
            base_addr = base_info["base_addr"]

            asm_name = self.asm_output_name(idx, item, bin_path, role_class, base_addr)
            asm_path = self.output_dir / asm_name

            cmd = build_objdump_command(
                bin_filename=bin_path.name,
                base_addr=base_addr,
                objdump=self.args.objdump,
                use_wsl=not self.args.no_wsl,
                extra_args=self.args.objdump_extra or [],
            )

            rec = self.record_base(item, idx, role_class)
            rec.update({
                "status": "PENDING",
                "input_rel": rel,
                "input_sha256": actual_sha,
                "input_size": len(data),
                "asm_rel": relpath_str(asm_path, self.workspace),
                "base_addr": base_addr,
                "base_detection": base_info,
                "command": cmd,
                "notes": notes,
            })

            if self.args.dry_run:
                self.counts["dry_run_records"] += 1
                rec["status"] = "DRY_RUN"
                rec["message"] = "command recorded but not executed"
                self.records.append(rec)
                self.add_inventory(rec)
                self.log(f"[DRY] {rel} -> {rec['asm_rel']} base={base_addr}")
                return

            self.counts["disassembly_attempts"] += 1
            asm_path.parent.mkdir(parents=True, exist_ok=True)
            with asm_path.open("w", encoding="utf-8", errors="replace") as out_f:
                proc = subprocess.run(
                    cmd,
                    cwd=str(bin_path.parent),
                    stdout=out_f,
                    stderr=subprocess.PIPE,
                    text=True,
                    check=False,
                )

            rec["returncode"] = proc.returncode
            rec["stderr"] = proc.stderr.strip() if proc.stderr else ""

            if proc.returncode == 0:
                self.counts["disassembly_succeeded"] += 1
                rec["status"] = "DISASSEMBLED"
                rec["message"] = "objdump completed"
                rec["asm_sha256"] = sha256_bytes(asm_path.read_bytes()) if asm_path.exists() else ""
                self.records.append(rec)
                self.add_inventory(rec)
                self.analyzer_inputs.append(self.to_analyzer_input(rec))
                self.log(f"[INFO] disassembled {rel} -> {rec['asm_rel']} base={base_addr}")
            else:
                self.counts["disassembly_failed"] += 1
                rec["status"] = "FAILED"
                rec["message"] = "objdump failed"
                self.records.append(rec)
                self.add_inventory(rec)
                self.log(f"[ERROR] objdump failed for {rel}: {rec['stderr']}")

        except Exception as e:
            self.counts["errors"] += 1
            rec = self.record_base(item, idx, role_class)
            rec.update({
                "status": "ERROR",
                "message": f"{type(e).__name__}: {e}",
            })
            self.records.append(rec)
            self.log(f"[ERROR] failed {rel}: {type(e).__name__}: {e}")

    def choose_base_address(self, item: Dict[str, Any], role_class: str, bin_path: Path) -> Dict[str, Any]:
        if self.args.force_base:
            return {
                "base_addr": self.args.force_base,
                "method": "forced_cli",
                "confidence": "forced",
            }

        if role_class == "stub":
            return {
                "base_addr": self.args.stub_base,
                "method": "fixed_stub_base",
                "confidence": "high",
            }

        if role_class == "firmware":
            return detect_base_vma_programmatically(bin_path)

        return {
            "base_addr": self.args.unknown_base,
            "method": "default_unknown_base",
            "confidence": "low",
        }

    def asm_output_name(self, idx: int, item: Dict[str, Any], bin_path: Path, role_class: str, base_addr: str) -> str:
        source = item.get("source", {}) or {}
        ident = source.get("identity_hint", {}) or {}
        model = sanitize_name(str(ident.get("model", "") or "UNKNOWN"))
        version = sanitize_name(str(ident.get("version", "") or "UNK"))
        stem = sanitize_name(bin_path.stem)
        base_clean = base_addr.replace("0x", "").replace("0X", "")
        return f"{idx:04d}_{model}_{version}_{role_class}_{base_clean}_{stem}.asm"

    def record_base(self, item: Dict[str, Any], idx: int, role_class: str) -> Dict[str, Any]:
        return {
            "record_id": f"disasm_{idx:04d}",
            "record_type": "stage5_disassembly",
            "input_index": idx,
            "input_kind": item.get("input_kind"),
            "role": effective_role_from_item(item, role_class),
            "original_role": original_role_from_item(item),
            "effective_role": effective_role_from_item(item, role_class),
            "role_class": role_class,
            "target_scope": target_scope_from_item(item),
            "container_type": item.get("container_type"),
            "source_stage": item.get("source_stage"),
            "source": item.get("source", {}),
            "stage4_decrypt": item.get("decrypt", {}),
            "stage3_route": item.get("stage3_route", {}),
        }

    def to_analyzer_input(self, rec: Dict[str, Any]) -> Dict[str, Any]:
        return {
            "input_kind": "disassembled_payload",
            "role": rec.get("effective_role") or (rec.get("role") if rec.get("role") in ("firmware", "stub") else rec.get("role_class")),
            "original_role": rec.get("original_role"),
            "effective_role": rec.get("effective_role") or rec.get("role_class"),
            "role_class": rec.get("role_class"),
            "target_scope": rec.get("target_scope"),
            "container_type": rec.get("container_type"),
            "source_stage": "stage5_disassembly",
            "plaintext_rel": rec.get("input_rel"),
            "plaintext_sha256": rec.get("input_sha256"),
            "asm_rel": rec.get("asm_rel"),
            "asm_sha256": rec.get("asm_sha256"),
            "base_addr": rec.get("base_addr"),
            "base_detection": rec.get("base_detection"),
            "stage4_decrypt": rec.get("stage4_decrypt", {}),
            "stage3_route": rec.get("stage3_route", {}),
            "source": rec.get("source", {}),
        }

    def add_inventory(self, rec: Dict[str, Any]) -> None:
        bd = rec.get("base_detection", {}) or {}
        self.inventory_rows.append({
            "record_id": rec.get("record_id", ""),
            "status": rec.get("status", ""),
            "role": rec.get("role", ""),
            "original_role": rec.get("original_role", ""),
            "effective_role": rec.get("effective_role", ""),
            "role_class": rec.get("role_class", ""),
            "target_scope": rec.get("target_scope", ""),
            "container_type": rec.get("container_type", ""),
            "input_rel": rec.get("input_rel", ""),
            "asm_rel": rec.get("asm_rel", ""),
            "input_size": rec.get("input_size", ""),
            "input_sha256": rec.get("input_sha256", ""),
            "base_addr": rec.get("base_addr", ""),
            "base_method": bd.get("method", ""),
            "base_confidence": bd.get("confidence", ""),
            "count_40": bd.get("count_40", ""),
            "count_90": bd.get("count_90", ""),
            "returncode": rec.get("returncode", ""),
            "message": rec.get("message", ""),
        })

    def write_outputs(self) -> None:
        manifest = {
            "schema_version": 1,
            "producer": PRODUCER,
            "created_utc": utc_iso(),
            "workspace": str(self.workspace),
            "inputs_rel": relpath_str(self.inputs_path, self.workspace),
            "output_dir": relpath_str(self.output_dir, self.workspace),
            "log_rel": relpath_str(self.log_path, self.workspace) if self.log_path else "",
            "counts": self.counts,
            "stage_boundaries": {
                "stage5a_does": [
                    "disassemble_plaintext_firmware_and_stubs",
                    "detect_firmware_base_vma",
                    "assign_fixed_stub_base",
                    "write_disassembly_manifest",
                    "prepare_stage5b_analyzer_inputs",
                    "preserve_target_scope_metadata",
                    "preserve_original_role_and_effective_role_metadata",
                    "classify_validation_reference_as_firmware",
                ],
                "stage5a_does_not": [
                    "decrypt",
                    "checksum_verification",
                    "command_handler_analysis",
                    "flash_routine_analysis",
                    "profile_generation",
                ],
            },
            "tooling": {
                "objdump": self.args.objdump,
                "use_wsl": not self.args.no_wsl,
                "architecture": "mn10300",
                "endianness": "big",
            },
            "records": self.records,
        }

        manifest_path = self.workspace / "stage5_disassembly_manifest.json"
        manifest_path.write_text(json.dumps(manifest, indent=4), encoding="utf-8")
        self.log(f"[INFO] wrote Stage 5 disassembly manifest: {relpath_str(manifest_path, self.workspace)}")

        inv_path = self.workspace / "stage5_disassembly_inventory.csv"
        fields = [
            "record_id", "status", "role", "original_role", "effective_role", "role_class", "target_scope", "container_type",
            "input_rel", "asm_rel", "input_size", "input_sha256",
            "base_addr", "base_method", "base_confidence", "count_40", "count_90",
            "returncode", "message",
        ]
        with inv_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(f"[INFO] wrote Stage 5 disassembly inventory: {relpath_str(inv_path, self.workspace)}")

        # If dry-run, analyzer inputs are not useful yet. But include what would be next.
        if self.args.dry_run:
            analyzer_inputs = [
                self.to_analyzer_input(rec)
                for rec in self.records
                if rec.get("status") == "DRY_RUN"
            ]
        else:
            analyzer_inputs = self.analyzer_inputs

        analyzer_doc = {
            "schema_version": 1,
            "producer": PRODUCER,
            "created_utc": utc_iso(),
            "workspace": str(self.workspace),
            "source_manifest_rel": "stage5_disassembly_manifest.json",
            "analysis_input_count": len(analyzer_inputs),
            "analysis_inputs": analyzer_inputs,
            "review_inputs": self.review_inputs,
        }
        analyzer_path = self.workspace / "stage5b_analyzer_inputs.json"
        analyzer_path.write_text(json.dumps(analyzer_doc, indent=4), encoding="utf-8")
        self.log(f"[INFO] wrote Stage 5B analyzer inputs: {relpath_str(analyzer_path, self.workspace)}")

    def print_summary(self) -> None:
        print()
        print("=" * 72)
        print("HLDS STAGE 5A DISASSEMBLER SUMMARY")
        print("=" * 72)
        print(f"Workspace              : {self.workspace}")
        print(f"Inputs                 : {self.inputs_path}")
        print(f"Output dir             : {self.output_dir}")
        print(f"Manifest               : {self.workspace / 'stage5_disassembly_manifest.json'}")
        print(f"Inventory CSV          : {self.workspace / 'stage5_disassembly_inventory.csv'}")
        print(f"Stage5B inputs         : {self.workspace / 'stage5b_analyzer_inputs.json'}")
        print(f"Log file               : {self.log_path}")
        print("-" * 72)
        print(f"Analysis inputs seen   : {self.counts['analysis_inputs_seen']}")
        print(f"Review inputs seen     : {self.counts['review_inputs_seen']}")
        print(f"Inputs selected        : {self.counts['inputs_selected']}")
        print(f"Firmware inputs        : {self.counts['firmware_inputs']}")
        print(f"Stub inputs            : {self.counts['stub_inputs']}")
        print(f"Unknown inputs         : {self.counts['unknown_inputs']}")
        print(f"Orig unknown->firmware : {self.counts['original_role_unknown_promoted_to_firmware']}")
        print(f"Attempts               : {self.counts['disassembly_attempts']}")
        print(f"Succeeded              : {self.counts['disassembly_succeeded']}")
        print(f"Failed                 : {self.counts['disassembly_failed']}")
        print(f"Dry-run records        : {self.counts['dry_run_records']}")
        print(f"Missing files          : {self.counts['missing_input_files']}")
        print(f"SHA mismatches         : {self.counts['sha256_mismatches']}")
        print(f"Errors                 : {self.counts['errors']}")
        print("=" * 72)


def build_arg_parser() -> argparse.ArgumentParser:
    ap = argparse.ArgumentParser(
        description="Stage 5A HLDS/LG MN103S disassembler. Disassembles Stage 5 plaintext inputs."
    )
    ap.add_argument(
        "--workspace",
        default=r"D:\vscode-workspace\xbox-dvd-tools-v2",
        help="Workspace root containing stage5_analysis_inputs.json and plaintext payloads.",
    )
    ap.add_argument(
        "--inputs",
        default=None,
        help="Path to stage5_analysis_inputs.json. Defaults to <workspace>\\stage5_analysis_inputs.json.",
    )
    ap.add_argument(
        "--output-dir",
        default=None,
        help="Directory for .asm outputs. Defaults to <workspace>\\disassembly.",
    )
    ap.add_argument(
        "--objdump",
        default="mn10300-elf-objdump",
        help="Objdump executable name/path as visible inside WSL or local environment.",
    )
    ap.add_argument(
        "--no-wsl",
        action="store_true",
        help="Call objdump directly instead of via WSL.",
    )
    ap.add_argument(
        "--objdump-extra",
        action="append",
        default=[],
        help="Additional argument passed to objdump. Can be repeated.",
    )
    ap.add_argument(
        "--stub-base",
        default="0x40008000",
        help="Base address for stubs. Default: 0x40008000.",
    )
    ap.add_argument(
        "--unknown-base",
        default="0x40000000",
        help="Base address for unknown/review payloads when included. Default: 0x40000000.",
    )
    ap.add_argument(
        "--force-base",
        default=None,
        help="Force all inputs to this base address, overriding firmware/stub detection.",
    )
    ap.add_argument(
        "--roles",
        nargs="*",
        choices=["firmware", "stub", "unknown"],
        help="Optional role filter. Example: --roles firmware stub",
    )
    ap.add_argument(
        "--include-review",
        action="store_true",
        help="Also attempt review_inputs from stage5_analysis_inputs.json. Default is analysis_inputs only.",
    )
    ap.add_argument(
        "--dry-run",
        action="store_true",
        help="Record commands and outputs without invoking objdump.",
    )
    ap.add_argument(
        "--clean",
        action="store_true",
        help="Delete output disassembly directory before running.",
    )
    return ap


def main(argv: Optional[Sequence[str]] = None) -> int:
    args = build_arg_parser().parse_args(argv)
    s = Stage5Disassembler(args)
    try:
        s.run()
        if s.counts.get("errors", 0) or s.counts.get("missing_input_files", 0) or s.counts.get("disassembly_failed", 0):
            return 2
        return 0
    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())
