#!/usr/bin/env python3
r"""
hlds_stage1_extractor.py

Stage 1 HLDS/LG firmware payload extractor/stager. Revision 3.

This script only discovers, extracts, copies, hashes, deduplicates, and indexes
candidate .dld/.bin payloads. It does not split DLDs, decrypt, validate checksums,
classify firmware families, update fleet_profiles.json, or write into firmwarehq_LG.

JSON is the machine-readable handoff for Stage 2.
Log file is the human-readable chronological debug trail.

Default tool paths:
    --deark D:\vscode-workspace\xbox-dvd-tools-v2\bin\deark.exe
    --upx   D:\vscode-workspace\xbox-dvd-tools-v2\bin\upx.exe
    --7z    C:\Program Files\7-Zip\7z.exe
"""

from __future__ import annotations

import argparse
import csv
import hashlib
import json
import os
import re
import shutil
import struct
import subprocess
import tempfile
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple

PRODUCER = "hlds_stage1_extractor_r3"
SCHEMA_VERSION = 1

HLDS_MAGICS = (b"HL-DT-ST", b"HL-DL-ST")
DLD_HEADER_SIZE = 0x40

ARCHIVE_EXTS = {".zip", ".7z", ".rar", ".cab"}
EXE_EXTS = {".exe"}
PAYLOAD_EXTS = {".dld", ".bin"}

DEFAULT_MIN_CARVED_BIN_SIZE = 0x1000
DEFAULT_MAX_CARVED_BIN_SIZE = 0x1000000


def utc_stamp() -> str:
    return time.strftime("%Y%m%d_%H%M%S", time.gmtime())


def utc_iso() -> str:
    return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())


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


def rel(path: Path, root: Path) -> str:
    try:
        return str(path.resolve().relative_to(root.resolve())).replace("/", "\\")
    except Exception:
        return str(path)


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 sha256_bytes(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()


def read_u32_le(data: bytes, off: int) -> Optional[int]:
    if off < 0 or off + 4 > len(data):
        return None
    return struct.unpack("<I", data[off:off + 4])[0]


def read_u32_be(data: bytes, off: int) -> Optional[int]:
    if off < 0 or off + 4 > len(data):
        return None
    return struct.unpack(">I", data[off:off + 4])[0]


def run_silent(cmd: Sequence[str], cwd: Optional[Path] = None) -> subprocess.CompletedProcess:
    return subprocess.run(
        list(map(str, cmd)),
        cwd=str(cwd) if cwd else None,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
        errors="replace",
    )


def tool_exists(path_or_name: str) -> bool:
    if not path_or_name:
        return False
    return Path(path_or_name).exists() or shutil.which(path_or_name) is not None


def resolve_tool(path_or_name: str, fallback_name: str) -> str:
    if path_or_name and (Path(path_or_name).exists() or shutil.which(path_or_name)):
        return path_or_name
    found = shutil.which(fallback_name)
    return found or path_or_name or fallback_name


def is_inside(path: Path, root: Path) -> bool:
    try:
        path.resolve().relative_to(root.resolve())
        return True
    except Exception:
        return False


def short_method(method: str) -> str:
    method = method.lower().strip()
    aliases = {
        "direct_copy": "direct",
        "archive_extract": "7z",
        "exe_archive_extract": "7zexe",
        "upx_unpack": "upx",
        "diskexpress": "dxp",
        "embedded_dld_carve": "dldcarve",
        "pe_rcdata": "rcdata",
        "pe_overlay": "overlay",
        "swapped_zip": "swzip",
    }
    return aliases.get(method, sanitize(method, "method").lower())


def is_pe_with_nt_headers(path: Path) -> bool:
    try:
        with path.open("rb") as f:
            head = f.read(0x1000)
        if len(head) < 0x40 or head[:2] != b"MZ":
            return False
        pe_off = read_u32_le(head, 0x3C)
        if pe_off is None or pe_off < 0x40:
            return False
        if pe_off + 4 <= len(head):
            return head[pe_off:pe_off + 4] == b"PE\x00\x00"
        with path.open("rb") as f:
            f.seek(pe_off)
            return f.read(4) == b"PE\x00\x00"
    except Exception:
        return False


def dld_container_spans(data: bytes, min_payload: int = 0x1000, max_payload: int = 0x1000000) -> List[Tuple[int, int, Dict[str, Any]]]:
    """
    Locate likely HLDS DLD containers.
    This is only used to carve whole DLD blobs; it does not split or validate firmware.
    """
    spans: List[Tuple[int, int, Dict[str, Any]]] = []
    seen: set[Tuple[int, int]] = set()

    for magic in HLDS_MAGICS:
        pos = 0
        while True:
            idx = data.find(magic, pos)
            if idx < 0:
                break
            start = idx - 8
            if start >= 0 and start + DLD_HEADER_SIZE <= len(data):
                header = data[start:start + DLD_HEADER_SIZE]
                fmt = read_u32_le(header, 0x04)
                fw_plus_header = read_u32_be(header, 0x10)
                payload_size = read_u32_be(header, 0x1C)
                if (
                    fmt == 1
                    and fw_plus_header is not None
                    and payload_size is not None
                    and DLD_HEADER_SIZE < fw_plus_header <= payload_size + DLD_HEADER_SIZE
                    and min_payload <= payload_size <= max_payload
                ):
                    end = start + DLD_HEADER_SIZE + payload_size
                    if end <= len(data) and (start, end) not in seen:
                        seen.add((start, end))
                        ident = header[0x20:0x40].decode("ascii", errors="ignore").strip()
                        spans.append((start, end, {
                            "magic": magic.decode("ascii"),
                            "identity_string": ident,
                            "payload_size": payload_size,
                        }))
            pos = idx + 1

    spans.sort(key=lambda x: (x[0], -(x[1] - x[0])))
    return spans


def identity_from_dld_header_bytes(data: bytes) -> Optional[Tuple[str, str]]:
    spans = dld_container_spans(data[:min(len(data), 0x200000)])
    if not spans:
        return None
    _, _, info = spans[0]
    ident = info.get("identity_string", "")
    # Examples:
    #   RW/DVD GCC-4246N0X06  I
    #   DVD-ROM GDR8081N0110  L
    m = re.search(r"(GCC-?\d{4}[A-Z]|GDR-?\d{4}[A-Z]|GCC\d{4}[A-Z]|GDR\d{4}[A-Z])([A-Z0-9.]{3,6})", ident, re.I)
    if not m:
        return None
    model = m.group(1).upper()
    if model.startswith(("GCC", "GDR")) and "-" not in model[:4]:
        model = model[:3] + "-" + model[3:]
    return sanitize(model), sanitize(m.group(2).upper())


def version_from_names(model: str, paths: Iterable[Path]) -> Tuple[str, str]:
    """
    Best-effort version hint. Stage 2 must treat this as a hint, not truth.
    """
    model_clean = model.upper().replace("-", "")
    candidates: List[Tuple[int, str, str]] = []

    token_re = re.compile(
        r"(?<![A-Z0-9])("
        r"N?\d\.\d{2}(?:[-_][A-Z0-9]{1,3})?|"   # 1.01, N1.00-G4
        r"[A-Z]{1,3}\d{2}(?:[-_]\d{2})?|"        # YE05, EL02-06, CT01
        r"[A-Z]\d{3}|"                            # A101
        r"\d[A-Z]\d{2}|"                          # 0X06, 0J06
        r"\d{4}"                                  # 0110
        r")(?![A-Z0-9])",
        re.I,
    )

    for path in paths:
        parts = list(path.parts[-5:])
        for priority, text in enumerate(reversed(parts)):
            stem = Path(text).stem.upper()
            # Remove obvious model text to avoid model/version concatenation.
            stem_no_model = stem.replace(model.upper(), "").replace(model_clean, "")
            stem_no_model = stem_no_model.replace("FIRMWARE", "").replace("UPDATE", "")
            for m in token_re.finditer(stem_no_model):
                token = m.group(1).replace("_", "-").upper()
                if token in {"WIN", "DOS", "EW", "KW", "EXE", "ZIP"}:
                    continue
                score = 100 - priority
                # Prefer tokens in the filename over parent dirs.
                if text == path.name:
                    score += 20
                candidates.append((score, token, "filename_or_directory"))

    if candidates:
        candidates.sort(key=lambda x: (-x[0], len(x[1])))
        return sanitize(candidates[0][1]), candidates[0][2]

    return "UNK", "unknown"


@dataclass
class SourceRecord:
    root_file: Path
    source_path: Path
    method: str
    model_hint: str
    version_hint: str
    version_source: str
    detail: str = ""


class HumanLogger:
    def __init__(self, log_path: Path, verbose: bool) -> None:
        self.log_path = log_path
        self.verbose = verbose
        self.log_path.parent.mkdir(parents=True, exist_ok=True)
        self._fh = self.log_path.open("w", encoding="utf-8", errors="replace")

    def close(self) -> None:
        self._fh.close()

    def _emit(self, level: str, msg: str, terminal: bool = True) -> None:
        line = f"[{time.strftime('%H:%M:%S')}] [{level}] {msg}"
        self._fh.write(line + "\n")
        self._fh.flush()
        if terminal:
            print(line)

    def info(self, msg: str) -> None:
        self._emit("INFO", msg)

    def error(self, msg: str) -> None:
        self._emit("ERROR", msg)

    def debug(self, msg: str) -> None:
        if self.verbose:
            self._emit("DEBUG", msg)
        else:
            self._fh.write(f"[{time.strftime('%H:%M:%S')}] [DEBUG] {msg}\n")
            self._fh.flush()


class Stage1Extractor:
    def __init__(self, args: argparse.Namespace) -> None:
        self.workspace = Path(args.workspace).resolve()
        self.input_root = Path(args.input_root).resolve()
        self.clean = args.clean
        self.verbose = args.verbose
        self.max_depth = args.max_depth
        self.min_carved_bin_size = args.min_carved_bin_size
        self.max_carved_bin_size = args.max_carved_bin_size
        self.stage_newstyle_probes = args.stage_newstyle_probes

        self.seven_zip = resolve_tool(args.seven_zip, "7z")
        self.deark = resolve_tool(args.deark, "deark")
        self.upx = resolve_tool(args.upx, "upx")

        self.staging_dir = self.workspace / "staging_candidates"
        self.prm_dir = self.workspace / "prm"
        self.work_dir = self.workspace / "extractor_work"
        self.logs_dir = self.workspace / "logs"

        if self.clean:
            for d in (self.staging_dir, self.prm_dir, self.work_dir, self.logs_dir):
                if d.exists():
                    shutil.rmtree(d)

        self.staging_dir.mkdir(parents=True, exist_ok=True)
        self.prm_dir.mkdir(parents=True, exist_ok=True)
        self.work_dir.mkdir(parents=True, exist_ok=True)
        self.logs_dir.mkdir(parents=True, exist_ok=True)

        self.log = HumanLogger(self.logs_dir / f"hlds_stage1_extractor_{utc_stamp()}.log", self.verbose)

        self.payload_by_hash: Dict[str, Dict[str, Any]] = {}
        self.payload_order: List[str] = []
        self.visited_source_paths: set[str] = set()
        self.extracted_archive_hashes: set[str] = set()
        self.prm_by_hash: Dict[str, Dict[str, Any]] = {}

        self.counts: Dict[str, int] = {
            "files_scanned": 0,
            "archives_extracted": 0,
            "upx_unpacked": 0,
            "diskexpress_extracted": 0,
            "embedded_dld_carved": 0,
            "pe_rcdata_extracted": 0,
            "pe_overlay_extracted": 0,
            "payloads_staged": 0,
            "prm_copied": 0,
            "duplicates": 0,
            "errors": 0,
        }

    def close(self) -> None:
        self.log.close()

    def rel_ws(self, path: Path) -> str:
        return rel(path, self.workspace)

    def rel_input(self, path: Path) -> str:
        return rel(path, self.input_root)

    def is_own_output_path(self, path: Path) -> bool:
        return (
            is_inside(path, self.staging_dir)
            or is_inside(path, self.work_dir)
            or is_inside(path, self.logs_dir)
        )

    def model_from_root_file(self, root_file: Path) -> Tuple[str, str]:
        try:
            rp = root_file.resolve().relative_to(self.input_root)
            if len(rp.parts) >= 2:
                return sanitize(rp.parts[0]), "top_level_directory"
        except Exception:
            pass
        return sanitize(root_file.parent.name), "parent_directory"

    def infer_identity(self, ext: str, data: bytes, source: SourceRecord) -> Dict[str, str]:
        model = source.model_hint
        model_source = "top_level_directory"

        if ext.lower() == ".dld":
            dld_id = identity_from_dld_header_bytes(data)
            if dld_id:
                dld_model, dld_ver = dld_id
                # Prefer top-level folder for model unless DLD has no matching folder context.
                if model.upper() in {"UNK", "UNKNOWN"}:
                    model = dld_model
                    model_source = "dld_header"
                return {
                    "model": sanitize(model),
                    "version": sanitize(dld_ver),
                    "model_source": model_source,
                    "version_source": "dld_header",
                }

        version, version_source = version_from_names(model, [source.source_path, source.root_file])
        return {
            "model": sanitize(model),
            "version": sanitize(version),
            "model_source": model_source,
            "version_source": version_source,
        }

    def staged_name(self, identity: Dict[str, str], method: str, sha: str, ext: str) -> str:
        model = sanitize(identity.get("model", "UNK"))
        version = sanitize(identity.get("version", "UNK"))
        method_s = short_method(method)
        return f"{model}_{version}_{method_s}_{sha[:8]}{ext.lower()}"

    def source_json(self, src: SourceRecord) -> Dict[str, str]:
        row = {
            "root_file_rel": self.rel_input(src.root_file) if is_inside(src.root_file, self.input_root) else self.rel_ws(src.root_file),
            "source_path_rel": self.rel_input(src.source_path) if is_inside(src.source_path, self.input_root) else self.rel_ws(src.source_path),
            "method": short_method(src.method),
        }
        if src.detail:
            row["detail"] = src.detail
        return row

    def stage_payload_bytes(self, data: bytes, ext: str, src: SourceRecord) -> None:
        ext = ext.lower()
        if ext not in PAYLOAD_EXTS:
            raise ValueError(f"invalid payload ext: {ext}")

        sha = sha256_bytes(data)
        size = len(data)
        if sha in self.payload_by_hash:
            self.counts["duplicates"] += 1
            rec = self.payload_by_hash[sha]
            rec["sources"].append(self.source_json(src))
            self.log.info(f"duplicate {ext}: {src.source_path.name} -> {rec['staged_rel']}")
            return

        identity = self.infer_identity(ext, data, src)
        staged = self.staging_dir / self.staged_name(identity, src.method, sha, ext)

        # Extremely unlikely name collision, but keep the first-hash rule intact.
        if staged.exists():
            staged = self.staging_dir / f"{staged.stem}_{sha[8:12]}{ext}"

        staged.write_bytes(data)

        payload_id = sha[:16]
        rec = {
            "payload_id": payload_id,
            "ext": ext,
            "staged_rel": self.rel_ws(staged),
            "size": size,
            "sha256": sha,
            "identity_hint": identity,
            "first_source": self.source_json(src),
            "sources": [self.source_json(src)],
        }

        self.payload_by_hash[sha] = rec
        self.payload_order.append(sha)
        self.counts["payloads_staged"] += 1
        self.log.info(f"staged {ext}: {staged.name} ({size} bytes)")

    def stage_payload_file(self, path: Path, src: SourceRecord, ext_override: Optional[str] = None) -> None:
        ext = (ext_override or path.suffix).lower()
        with path.open("rb") as f:
            data = f.read()
        self.stage_payload_bytes(data, ext, src)

    def run(self) -> None:
        self.log.info(f"stage1 start: input={self.input_root}")
        self.log.info(f"workspace={self.workspace}")
        self.log.info(f"tools: 7z={self.seven_zip} deark={self.deark} upx={self.upx}")

        if not self.input_root.exists():
            raise FileNotFoundError(self.input_root)

        if self.input_root.is_file():
            model_hint, _ = self.model_from_root_file(self.input_root)
            self.process_file(self.input_root, self.input_root, model_hint, depth=0)
        else:
            for path in sorted(self.input_root.rglob("*")):
                if path.is_file():
                    if self.is_own_output_path(path):
                        continue
                    model_hint, _ = self.model_from_root_file(path)
                    self.process_file(path, path, model_hint, depth=0)

        self.write_outputs()
        self.log.info("stage1 complete")
        self.print_summary()

    def process_file(self, path: Path, root_file: Path, model_hint: str, depth: int, allow_workdir: bool = False) -> None:
        if depth > self.max_depth:
            self.log.debug(f"max depth reached: {path}")
            return
        if not path.exists() or not path.is_file():
            return

        # Never process staged output, PRM output, or logs.
        # Files under extractor_work are allowed only when they are produced by
        # the current recursive extraction path. r2 incorrectly skipped these,
        # which caused archive-extracted .dld/.bin/.prm files to be missed.
        if is_inside(path, self.staging_dir) or is_inside(path, self.prm_dir) or is_inside(path, self.logs_dir):
            return
        if is_inside(path, self.work_dir) and not allow_workdir:
            return

        key = str(path.resolve()).lower()
        if key in self.visited_source_paths:
            return
        self.visited_source_paths.add(key)

        self.counts["files_scanned"] += 1
        suffix = path.suffix.lower()

        try:
            if suffix == ".prm":
                self.copy_prm(path, root_file, model_hint)
                return

            if suffix in PAYLOAD_EXTS:
                src = SourceRecord(
                    root_file=root_file,
                    source_path=path,
                    method="direct_copy" if is_inside(path, self.input_root) else "archive_extract",
                    model_hint=model_hint,
                    version_hint="",
                    version_source="",
                )
                self.stage_payload_file(path, src)
                return

            if suffix in ARCHIVE_EXTS:
                self.extract_archive(path, root_file, model_hint, depth, method="archive_extract")
                self.try_swapped_zip(path, root_file, model_hint, depth)
                return

            if suffix in EXE_EXTS:
                self.process_exe(path, root_file, model_hint, depth)
                return

            if self.is_dld_carve_candidate(path):
                self.carve_embedded_dld(path, root_file, model_hint)

            # Some extracted PE resource folders contain extensionless RCDATA files such as 103.
            if self.is_rcdata_like_path(path) and self.is_carved_bin_size(path.stat().st_size):
                if not self.stage_newstyle_probes:
                    self.log.info(f"newstyle probe extensionless RCDATA not staged: {path.name} size=0x{path.stat().st_size:X}")
                    return
                src = SourceRecord(
                    root_file=root_file,
                    source_path=path,
                    method=f"pe_rcdata_{path.name}",
                    model_hint=model_hint,
                    version_hint="",
                    version_source="",
                    detail="extensionless RCDATA staged as .bin",
                )
                self.stage_payload_file(path, src, ext_override=".bin")
                return

        except Exception as e:
            self.counts["errors"] += 1
            self.log.error(f"{path}: {type(e).__name__}: {e}")

    def is_dld_carve_candidate(self, path: Path) -> bool:
        """Wrapper-like files worth scanning for embedded whole DLD containers.

        We deliberately do not scan .bin payloads for DLD by default; .bin is already a
        staged payload candidate and scanning it creates confusing double records.
        """
        suffix = path.suffix.lower()
        if suffix in PAYLOAD_EXTS or suffix == ".prm":
            return False
        if suffix in {".txt", ".rtf", ".ini", ".inf", ".cat", ".bmp", ".ico", ".cur", ".dll", ".sys", ".bat", ".cmd"}:
            return False
        if suffix in {".exe", ".scr", ".com", ".dat", ".tmp", ".fls", ".rom", ".img", ".ima", ".dxp"}:
            return True
        # Extensionless files from archive/tool outputs are often raw wrappers.
        if suffix == "" and not is_inside(path, self.input_root):
            return True
        return False

    def copy_prm(self, path: Path, root_file: Path, model_hint: str) -> None:
        """Copy .prm files into workspace/prm for later side-project analysis."""
        try:
            digest = sha256_file(path)
            if digest in self.prm_by_hash:
                self.prm_by_hash[digest]["sources"].append({
                    "root_file_rel": self.rel_input(root_file) if is_inside(root_file, self.input_root) else self.rel_ws(root_file),
                    "source_path_rel": self.rel_input(path) if is_inside(path, self.input_root) else self.rel_ws(path),
                })
                self.log.info(f"duplicate .prm: {path.name} -> {self.prm_by_hash[digest]['copied_rel']}")
                return

            model = sanitize(model_hint)
            stem = sanitize(path.stem, "PRM")
            dest = self.prm_dir / f"{model}_{stem}_{digest[:8]}.prm"
            shutil.copy2(path, dest)
            row = {
                "copied_rel": self.rel_ws(dest),
                "sha256": digest,
                "size": path.stat().st_size,
                "model_hint": model,
                "first_source": self.rel_input(path) if is_inside(path, self.input_root) else self.rel_ws(path),
                "sources": [{
                    "root_file_rel": self.rel_input(root_file) if is_inside(root_file, self.input_root) else self.rel_ws(root_file),
                    "source_path_rel": self.rel_input(path) if is_inside(path, self.input_root) else self.rel_ws(path),
                }],
            }
            self.prm_by_hash[digest] = row
            self.counts["prm_copied"] += 1
            self.log.info(f"copied .prm: {path.name} -> {self.rel_ws(dest)}")
        except Exception as e:
            self.counts["errors"] += 1
            self.log.error(f"PRM copy failed {path}: {type(e).__name__}: {e}")

    def process_exe(self, path: Path, root_file: Path, model_hint: str, depth: int) -> None:
        self.try_upx(path, root_file, model_hint, depth)
        self.carve_embedded_dld(path, root_file, model_hint)
        self.try_diskexpress(path, root_file, model_hint)
        self.extract_pe_rcdata(path, root_file, model_hint)
        self.extract_pe_overlay(path, root_file, model_hint)

        # Many updater EXEs are also archive-like.
        self.extract_archive(path, root_file, model_hint, depth, method="exe_archive_extract", optional=True)
        self.try_swapped_zip(path, root_file, model_hint, depth)

    def extract_archive(
        self,
        path: Path,
        root_file: Path,
        model_hint: str,
        depth: int,
        method: str,
        optional: bool = False,
    ) -> None:
        if not tool_exists(self.seven_zip):
            if not optional:
                self.log.error(f"7z missing; cannot extract {path}")
                self.counts["errors"] += 1
            else:
                self.log.debug(f"7z missing; optional extract skipped: {path}")
            return

        try:
            digest = sha256_file(path)
        except Exception:
            digest = sanitize(path.name)

        # Avoid extracting the exact same archive contents multiple times.
        archive_key = f"{method}:{digest}"
        if archive_key in self.extracted_archive_hashes:
            self.log.debug(f"archive already extracted by hash: {path}")
            return
        self.extracted_archive_hashes.add(archive_key)

        out_dir = self.work_dir / f"{sanitize(path.stem)}_{digest[:8]}_{short_method(method)}"
        out_dir.mkdir(parents=True, exist_ok=True)

        result = run_silent([self.seven_zip, "x", str(path), f"-o{out_dir}", "-y", "-bd"])
        if result.returncode != 0:
            self.log.debug(f"{short_method(method)} failed/not archive: {path} rc={result.returncode}")
            return

        self.counts["archives_extracted"] += 1
        self.log.info(f"extracted archive: {path.name} -> {self.rel_ws(out_dir)}")

        for child in sorted(out_dir.rglob("*")):
            if child.is_file():
                self.process_file(child, root_file, model_hint, depth + 1, allow_workdir=True)

    def try_swapped_zip(self, path: Path, root_file: Path, model_hint: str, depth: int) -> None:
        """
        Simple swapped-byte ZIP recovery used for odd KP/PK-flipped payloads.
        This only writes into extractor_work.
        """
        try:
            data = path.read_bytes()
        except Exception:
            return
        if len(data) < 4:
            return

        swapped = bytearray(data)
        for i in range(0, len(swapped) - 1, 2):
            swapped[i], swapped[i + 1] = swapped[i + 1], swapped[i]

        if not (bytes(swapped[:4]) == b"PK\x03\x04"):
            self.log.debug(f"swapped zip not detected: {path.name}")
            return

        digest = sha256_bytes(swapped)
        sw_path = self.work_dir / f"{sanitize(path.stem)}_{digest[:8]}_swapped.zip"
        sw_path.write_bytes(swapped)
        self.log.info(f"recovered swapped zip: {path.name} -> {self.rel_ws(sw_path)}")
        self.extract_archive(sw_path, root_file, model_hint, depth + 1, method="swapped_zip")

    def try_upx(self, path: Path, root_file: Path, model_hint: str, depth: int) -> None:
        if not tool_exists(self.upx):
            self.log.debug(f"upx missing; skipped {path.name}")
            return

        try:
            digest = sha256_file(path)
        except Exception:
            digest = sanitize(path.name)
        out = self.work_dir / f"{sanitize(path.stem)}_{digest[:8]}_upx.exe"

        result = run_silent([self.upx, "-d", f"-o{out}", str(path)])
        if result.returncode != 0 or not out.exists() or out.stat().st_size <= 0:
            self.log.debug(f"upx not packed/failed: {path.name} rc={result.returncode}")
            if out.exists() and out.stat().st_size == 0:
                try:
                    out.unlink()
                except Exception:
                    pass
            return

        self.counts["upx_unpacked"] += 1
        self.log.info(f"upx unpacked: {path.name} -> {self.rel_ws(out)}")
        self.process_file(out, root_file, model_hint, depth + 1, allow_workdir=True)

    def carve_embedded_dld(self, path: Path, root_file: Path, model_hint: str) -> None:
        try:
            data = path.read_bytes()
        except Exception as e:
            self.counts["errors"] += 1
            self.log.error(f"read failed for DLD carve {path}: {e}")
            return

        spans = dld_container_spans(data)
        for idx, (start, end, info) in enumerate(spans):
            blob = data[start:end]
            src = SourceRecord(
                root_file=root_file,
                source_path=path,
                method="embedded_dld_carve",
                model_hint=model_hint,
                version_hint="",
                version_source="",
                detail=f"offset=0x{start:X}; size=0x{len(blob):X}; idx={idx}",
            )
            self.stage_payload_bytes(blob, ".dld", src)
            self.counts["embedded_dld_carved"] += 1

    def try_diskexpress(self, path: Path, root_file: Path, model_hint: str) -> None:
        if not tool_exists(self.deark):
            self.log.debug(f"deark missing; diskexpress skipped: {path.name}")
            return
        if not tool_exists(self.seven_zip):
            self.log.debug(f"7z missing; diskexpress image extraction skipped: {path.name}")
            return

        try:
            data = path.read_bytes()
        except Exception:
            return

        if b"AS\x02\x00" not in data:
            self.log.debug(f"diskexpress signature not found: {path.name}")
            return

        candidates: List[Tuple[str, bytes]] = []

        marker = b"AS\x02\x00"
        idx_direct = data.find(marker)
        if idx_direct >= 0:
            candidates.append((f"marker_0x{idx_direct:X}", data[idx_direct:]))

        if data[:2] == b"MZ" and len(data) >= 28:
            e_cblp = struct.unpack("<H", data[2:4])[0]
            e_cp = struct.unpack("<H", data[4:6])[0]
            stub_size = e_cp * 512 if e_cblp == 0 else ((e_cp - 1) * 512) + e_cblp
            if 0 < stub_size < len(data):
                stripped = data[stub_size:]
                idx = stripped.find(marker)
                if idx >= 0:
                    candidates.append((f"mz_stub_stripped_0x{stub_size:X}_marker_0x{idx:X}", stripped[idx:]))

        # Dedupe equivalent DXP streams.
        seen_streams: set[str] = set()
        found_total = 0

        for label, stream in candidates:
            digest = sha256_bytes(stream)
            if digest in seen_streams:
                continue
            seen_streams.add(digest)

            with tempfile.TemporaryDirectory(prefix="hlds_dxp_", dir=str(self.work_dir)) as td:
                td_path = Path(td)
                dxp_path = td_path / f"{sanitize(path.stem)}_{label}.dxp"
                dxp_path.write_bytes(stream)

                result = run_silent([self.deark, "-m", "dskexp", str(dxp_path)], cwd=td_path)
                if result.returncode != 0:
                    self.log.debug(f"deark diskexpress failed: {path.name} {label} rc={result.returncode}")
                    continue

                # Deark may create output files at top level or in nested dirs.
                image_candidates = []
                for p in td_path.rglob("*"):
                    if p.is_file() and p != dxp_path:
                        # Try anything image-like or non-trivially sized. 7z will reject non-images.
                        if p.suffix.lower() in {".ima", ".img", ".bin", ".raw", ""} or p.stat().st_size > 0x10000:
                            image_candidates.append(p)

                for image_idx, ima in enumerate(sorted(image_candidates)):
                    extract_dir = td_path / f"image_extract_{image_idx}"
                    extract_dir.mkdir(exist_ok=True)
                    run_silent([self.seven_zip, "x", str(ima), f"-o{extract_dir}", "-y", "-bd"])

                    for child in sorted(extract_dir.rglob("*")):
                        if child.is_file() and child.suffix.lower() in PAYLOAD_EXTS:
                            src = SourceRecord(
                                root_file=root_file,
                                source_path=path,
                                method="diskexpress",
                                model_hint=model_hint,
                                version_hint="",
                                version_source="",
                                detail=f"dxp={label}; image={ima.name}; extracted_name={child.name}",
                            )
                            self.stage_payload_file(child, src)
                            found_total += 1

                # Also stage any direct .dld/.bin Deark output, in case no image extraction is needed.
                for child in sorted(td_path.rglob("*")):
                    if child.is_file() and child.suffix.lower() in PAYLOAD_EXTS and child != dxp_path:
                        src = SourceRecord(
                            root_file=root_file,
                            source_path=path,
                            method="diskexpress",
                            model_hint=model_hint,
                            version_hint="",
                            version_source="",
                            detail=f"dxp={label}; direct_deark_output={child.name}",
                        )
                        self.stage_payload_file(child, src)
                        found_total += 1

        if found_total:
            self.counts["diskexpress_extracted"] += found_total
            self.log.info(f"diskexpress payloads: {path.name} -> {found_total}")

    def extract_pe_rcdata(self, path: Path, root_file: Path, model_hint: str) -> None:
        if not is_pe_with_nt_headers(path):
            self.log.debug(f"not PE32 for RCDATA: {path.name}")
            return

        try:
            import pefile  # type: ignore
        except Exception:
            self.log.debug("pefile module missing; PE RCDATA extraction skipped")
            return

        try:
            pe = pefile.PE(str(path))
        except Exception as e:
            self.log.debug(f"pefile parse skipped: {path.name}: {e}")
            return

        if not hasattr(pe, "DIRECTORY_ENTRY_RESOURCE"):
            self.log.debug(f"PE has no resources: {path.name}")
            return

        try:
            for resource_type in pe.DIRECTORY_ENTRY_RESOURCE.entries:
                # RT_RCDATA = 10
                if getattr(resource_type, "id", None) != 10:
                    continue

                for entry in resource_type.directory.entries:
                    try:
                        data_entry = entry.directory.entries[0].data.struct
                        rva = data_entry.OffsetToData
                        size = data_entry.Size
                        if not self.is_carved_bin_size(size):
                            self.log.debug(f"RCDATA skipped by size: {path.name} id={getattr(entry, 'id', 'unknown')} size={size}")
                            continue
                        blob = pe.get_memory_mapped_image()[rva:rva + size]
                        res_id = str(getattr(entry, "id", "unknown"))
                        if not self.stage_newstyle_probes:
                            self.log.info(f"newstyle probe RCDATA not staged: {path.name} id={res_id} size=0x{size:X}")
                            continue
                        src = SourceRecord(
                            root_file=root_file,
                            source_path=path,
                            method=f"pe_rcdata_{res_id}",
                            model_hint=model_hint,
                            version_hint="",
                            version_source="",
                            detail=f"resource_id={res_id}; rva=0x{rva:X}; size=0x{size:X}",
                        )
                        self.stage_payload_bytes(blob, ".bin", src)
                        self.counts["pe_rcdata_extracted"] += 1
                    except Exception as e:
                        self.counts["errors"] += 1
                        self.log.error(f"RCDATA extraction error {path.name}: {e}")
        finally:
            try:
                pe.close()
            except Exception:
                pass

    def extract_pe_overlay(self, path: Path, root_file: Path, model_hint: str) -> None:
        try:
            data = path.read_bytes()
        except Exception:
            return
        if data[:2] != b"MZ" or len(data) < 0x100:
            return
        pe_off = read_u32_le(data, 0x3C)
        if pe_off is None or pe_off + 0x18 >= len(data) or data[pe_off:pe_off + 4] != b"PE\x00\x00":
            self.log.debug(f"not PE32 for overlay: {path.name}")
            return

        try:
            num_sections = struct.unpack("<H", data[pe_off + 6:pe_off + 8])[0]
            opt_size = struct.unpack("<H", data[pe_off + 20:pe_off + 22])[0]
            sec_table = pe_off + 24 + opt_size
            sections: List[Tuple[int, int]] = []

            for i in range(num_sections):
                off = sec_table + i * 40
                if off + 40 > len(data):
                    break
                raw_size = read_u32_le(data, off + 16) or 0
                raw_ptr = read_u32_le(data, off + 20) or 0
                if raw_ptr and raw_size and raw_ptr + raw_size <= len(data):
                    sections.append((raw_ptr, raw_ptr + raw_size))

            if not sections:
                return

            sections.sort()
            candidates: List[Tuple[str, int, bytes]] = []

            last_end = max(e for _, e in sections)
            if self.is_carved_bin_size(len(data) - last_end):
                candidates.append(("pe_overlay", last_end, data[last_end:]))

            for i in range(len(sections) - 1):
                start = sections[i][1]
                end = sections[i + 1][0]
                if self.is_carved_bin_size(end - start):
                    candidates.append((f"pe_gap_{i}", start, data[start:end]))

            for method, off, blob in candidates:
                if not self.stage_newstyle_probes:
                    self.log.info(f"newstyle probe {method} not staged: {path.name} offset=0x{off:X} size=0x{len(blob):X}")
                    continue
                src = SourceRecord(
                    root_file=root_file,
                    source_path=path,
                    method=method,
                    model_hint=model_hint,
                    version_hint="",
                    version_source="",
                    detail=f"offset=0x{off:X}; size=0x{len(blob):X}",
                )
                self.stage_payload_bytes(blob, ".bin", src)
                self.counts["pe_overlay_extracted"] += 1
        except Exception as e:
            self.counts["errors"] += 1
            self.log.error(f"PE overlay extraction error {path.name}: {e}")

    def is_carved_bin_size(self, size: int) -> bool:
        return self.min_carved_bin_size <= size <= self.max_carved_bin_size

    def is_rcdata_like_path(self, path: Path) -> bool:
        parts = [p.upper() for p in path.parts]
        return "RCDATA" in parts or path.name.upper().startswith("RCDATA")

    def write_outputs(self) -> None:
        payloads = [self.payload_by_hash[sha] for sha in self.payload_order]

        index = {
            "schema_version": SCHEMA_VERSION,
            "producer": PRODUCER,
            "created_utc": utc_iso(),
            "workspace": str(self.workspace),
            "input_root": str(self.input_root),
            "staging_dir": self.rel_ws(self.staging_dir),
            "log_rel": self.rel_ws(self.log.log_path),
            "counts": self.counts,
            "payloads": payloads,
            "side_artifacts": {
                "prm": list(self.prm_by_hash.values()),
            },
        }

        prm_index_path = self.workspace / "prm_index.json"
        prm_index = {
            "schema_version": 1,
            "producer": PRODUCER,
            "created_utc": utc_iso(),
            "workspace": str(self.workspace),
            "prm_dir": self.rel_ws(self.prm_dir),
            "count": len(self.prm_by_hash),
            "items": list(self.prm_by_hash.values()),
        }
        prm_index_path.write_text(json.dumps(prm_index, indent=4), encoding="utf-8")
        self.log.info(f"wrote PRM index: {self.rel_ws(prm_index_path)}")

        index_path = self.workspace / "stage1_payload_index.json"
        index_path.write_text(json.dumps(index, indent=4), encoding="utf-8")
        self.log.info(f"wrote JSON index: {self.rel_ws(index_path)}")

        csv_path = self.workspace / "stage1_inventory.csv"
        with csv_path.open("w", newline="", encoding="utf-8") as f:
            fields = [
                "payload_id", "ext", "staged_rel", "size", "sha256",
                "model", "version", "version_source", "source_count", "first_method", "first_root_file",
            ]
            writer = csv.DictWriter(f, fieldnames=fields)
            writer.writeheader()
            for p in payloads:
                ih = p["identity_hint"]
                fs = p["first_source"]
                writer.writerow({
                    "payload_id": p["payload_id"],
                    "ext": p["ext"],
                    "staged_rel": p["staged_rel"],
                    "size": p["size"],
                    "sha256": p["sha256"],
                    "model": ih.get("model", ""),
                    "version": ih.get("version", ""),
                    "version_source": ih.get("version_source", ""),
                    "source_count": len(p.get("sources", [])),
                    "first_method": fs.get("method", ""),
                    "first_root_file": fs.get("root_file_rel", ""),
                })
        self.log.info(f"wrote CSV inventory: {self.rel_ws(csv_path)}")

    def print_summary(self) -> None:
        print("")
        print("=" * 72)
        print("HLDS STAGE 1 EXTRACTOR SUMMARY")
        print("=" * 72)
        print(f"Input root        : {self.input_root}")
        print(f"Staging dir       : {self.staging_dir}")
        print(f"PRM dir           : {self.prm_dir}")
        print(f"JSON index        : {self.workspace / 'stage1_payload_index.json'}")
        print(f"Inventory CSV     : {self.workspace / 'stage1_inventory.csv'}")
        print(f"Log file          : {self.log.log_path}")
        print("-" * 72)
        print(f"Files scanned     : {self.counts['files_scanned']}")
        print(f"Archives extracted: {self.counts['archives_extracted']}")
        print(f"UPX unpacked      : {self.counts['upx_unpacked']}")
        print(f"DiskExpress found : {self.counts['diskexpress_extracted']}")
        print(f"Embedded DLDs     : {self.counts['embedded_dld_carved']}")
        print(f"PE RCDATA bins    : {self.counts['pe_rcdata_extracted']}")
        print(f"PE overlay bins   : {self.counts['pe_overlay_extracted']}")
        print(f"Payloads staged   : {self.counts['payloads_staged']}")
        print(f"PRM copied        : {self.counts['prm_copied']}")
        print(f"Duplicates linked : {self.counts['duplicates']}")
        print(f"Errors            : {self.counts['errors']}")
        print("=" * 72)


def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace:
    ap = argparse.ArgumentParser(description="Stage 1 HLDS/LG firmware payload extractor/stager. Revision 3.")
    ap.add_argument("--workspace", default=r"D:\vscode-workspace\xbox-dvd-tools-v2", help="Workspace root for outputs.")
    ap.add_argument("--input-root", default=r"D:\vscode-workspace\xbox-dvd-tools-v2\firmwarehq_LG", help="Firmware directory to scan. Never modified.")
    ap.add_argument("--clean", action="store_true", help="Delete stage1 output directories before running.")
    ap.add_argument("--verbose", action="store_true", help="Print verbose/debug details to terminal as well as log.")
    ap.add_argument("--max-depth", type=int, default=8, help="Maximum recursive extraction depth.")
    ap.add_argument("--min-carved-bin-size", type=lambda x: int(x, 0), default=DEFAULT_MIN_CARVED_BIN_SIZE, help="Minimum size for carved/resource/overlay .bin candidates.")
    ap.add_argument("--max-carved-bin-size", type=lambda x: int(x, 0), default=DEFAULT_MAX_CARVED_BIN_SIZE, help="Maximum size for carved/resource/overlay .bin candidates.")
    ap.add_argument("--stage-newstyle-probes", action="store_true", help="Stage ambiguous PE RCDATA / overlay / gap .bin probes. Default is log-only so older DLD-era extraction stays clean.")

    ap.add_argument("--7z", dest="seven_zip", default=r"C:\Program Files\7-Zip\7z.exe", help="Path to 7z.exe.")
    ap.add_argument("--deark", default=r"D:\vscode-workspace\xbox-dvd-tools-v2\bin\deark.exe", help="Path to deark.exe; needed for Disk eXpress EXEs.")
    ap.add_argument("--upx", default=r"D:\vscode-workspace\xbox-dvd-tools-v2\bin\upx.exe", help="Path to upx.exe; optional.")
    return ap.parse_args(argv)


def main(argv: Optional[Sequence[str]] = None) -> int:
    args = parse_args(argv)
    extractor = Stage1Extractor(args)
    try:
        extractor.run()
        return 0 if extractor.counts["errors"] == 0 else 2
    finally:
        extractor.close()


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