#!/usr/bin/env python3
"""Discover firmware/emulator candidate files from a workspace tree.

This helper is deliberately workspace/file driven. It does not consume Stage6,
Stage6A, profile_data.zip, or generated model profiles. It exists only to make
it easier to build explicit emulator JSON configs for many different HLDS
models from the files already present in a repo.
"""
from __future__ import annotations

import argparse
import json
import os
import re
import sys
from collections import defaultdict
from pathlib import Path
from typing import Any

MODEL_RE = re.compile(r"(?:^|[/\\_\-])((?:GCC|GDR|GSA|GWA|GCE|GCA|GBW|GBC|GH|BH|WH|GGW|GGC|GT|GU|GP|GS|CT|BE|BP|CH|UH)[A-Z0-9\-()]+)", re.I)
ASM_BASE_RE = re.compile(r"^\s*([0-9A-Fa-f]{8})\s+<")
ADDR_IN_NAME_RE = re.compile(r"(?:^|_)([49][0-9A-Fa-f]{7})(?:_|$)")
FW_EXTS = {".bin", ".dld"}
ASM_EXTS = {".asm", ".s"}


def rel_to(path: Path, root: Path) -> str:
    try:
        return path.resolve().relative_to(root.resolve()).as_posix()
    except Exception:
        return path.as_posix()


def infer_model(path: Path, workspace: Path) -> str:
    parts = path.parts
    # Prefer the directory directly under firmware/ when present.
    for i, part in enumerate(parts):
        if part.lower() == "firmware" and i + 1 < len(parts):
            return parts[i + 1]
    text = path.as_posix()
    m = MODEL_RE.search(text)
    return m.group(1).upper() if m else "UNKNOWN"


def infer_base_from_name(path: Path) -> str | None:
    m = ADDR_IN_NAME_RE.search(path.name)
    if m:
        return "0x" + m.group(1).upper()
    return None


def infer_base_from_asm(path: Path) -> str | None:
    try:
        with path.open("r", encoding="utf-8", errors="replace") as fh:
            for _ in range(80):
                line = fh.readline()
                if not line:
                    break
                m = ASM_BASE_RE.match(line)
                if m:
                    return "0x" + m.group(1).upper()
    except OSError:
        return None
    return None


def classify_file(path: Path) -> str:
    name = path.name.lower()
    if "stub" in name or "prepstub" in name or "loader" in name:
        return "stub_or_loader"
    if "plaintext" in name or "plain" in name or path.suffix.lower() in {".bin", ".dld"}:
        return "firmware_or_payload"
    return "unknown"


def discover(workspace: Path, roots: list[str], models: set[str] | None) -> list[dict[str, Any]]:
    candidates: list[dict[str, Any]] = []
    for root_name in roots:
        root = workspace / root_name
        if not root.exists():
            continue
        for path in root.rglob("*"):
            if not path.is_file():
                continue
            ext = path.suffix.lower()
            if ext not in FW_EXTS and ext not in ASM_EXTS:
                continue
            model = infer_model(path, workspace)
            if models and model.upper() not in {m.upper() for m in models}:
                continue
            base = infer_base_from_name(path)
            if ext in ASM_EXTS:
                base = infer_base_from_asm(path) or base
            candidates.append({
                "model": model,
                "kind": "asm" if ext in ASM_EXTS else "binary",
                "role_guess": classify_file(path),
                "path_rel": rel_to(path, workspace),
                "base_addr_guess": base or "0x90000000",
                "size": path.stat().st_size,
            })
    return sorted(candidates, key=lambda r: (r["model"], r["kind"], r["path_rel"]))


def summarize(rows: list[dict[str, Any]]) -> dict[str, Any]:
    by_model: dict[str, dict[str, int]] = defaultdict(lambda: defaultdict(int))
    for r in rows:
        by_model[r["model"]][r["kind"]] += 1
        by_model[r["model"]][r["role_guess"]] += 1
    return {
        "candidate_count": len(rows),
        "model_count": len(by_model),
        "models": {k: dict(v) for k, v in sorted(by_model.items())},
    }


def main(argv: list[str]) -> int:
    ap = argparse.ArgumentParser(description="Discover firmware candidates from workspace files, without profiles")
    ap.add_argument("--workspace", type=Path, default=Path.cwd())
    ap.add_argument("--roots", default="firmware,decrypted_payloads,stage5_disassembly", help="comma-separated workspace subtrees to scan")
    ap.add_argument("--models", help="comma-separated model filter, e.g. GCC-4241N,GDR-8050L")
    ap.add_argument("--summary", action="store_true")
    ap.add_argument("--out", type=Path, default=Path("firmware_target_discovery.json"))
    ns = ap.parse_args(argv[1:])
    roots = [x.strip() for x in ns.roots.split(",") if x.strip()]
    models = {x.strip() for x in ns.models.split(",")} if ns.models else None
    rows = discover(ns.workspace, roots, models)
    result = {"schema_version": 1, "source": "workspace_files_not_profiles", "workspace": str(ns.workspace), "summary": summarize(rows), "candidates": rows}
    ns.out.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    if ns.summary:
        print(json.dumps(result["summary"], indent=2, sort_keys=True))
    print(f"[OK] discovered {len(rows)} candidate file(s) -> {ns.out}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv))
