#!/usr/bin/env python3
"""Audit opcode/disassembly coverage from workspace .asm files.

This is a generic filesystem auditor. It does not load model profiles or
profile_data.zip. Use it to prioritize missing opcode work across whatever
firmware assembly listings are present in the workspace.
"""
from __future__ import annotations

import argparse
import json
import os
import re
import sys
from collections import Counter, 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)
INSN_RE = re.compile(r"^\s*([0-9a-fA-F]+):\s+((?:[0-9a-fA-F]{2}\s+){1,8})\s+([A-Za-z0-9_.]+)", re.I)
UNKNOWN_RE = re.compile(r"\bunknown\b", re.I)
UDF_RE = re.compile(r"\b(?:udf|break)\b", re.I)


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) -> str:
    parts = path.parts
    for i, part in enumerate(parts):
        if part.lower() == "firmware" and i + 1 < len(parts):
            return parts[i + 1]
    m = MODEL_RE.search(path.as_posix())
    return m.group(1).upper() if m else "UNKNOWN"


def classify_bytes(hex_bytes: str) -> str:
    bs = bytes.fromhex(" ".join(hex_bytes.split()))
    if not bs:
        return "empty"
    op = bs[0]
    if op == 0xFF:
        return "break_or_padding"
    # IDA's MN103 module often prints FE 00/01/02 and FE 80/81/82 bit
    # operations as unknown even though mn103s_emu.py implements them. Keep
    # them separate from true opcode gaps so completion work targets only
    # unimplemented or model-specific extension forms.
    if op == 0xFE and len(bs) >= 2 and bs[1] in {0x00, 0x01, 0x02, 0x80, 0x81, 0x82}:
        return "implemented_fe_bit_operation"
    if op in {0xF5, 0xF6, 0xF7, 0xF9, 0xFB, 0xFD}:
        return "custom_or_model_extension_space"
    if op in {0xF0, 0xF1, 0xF2, 0xF3, 0xF4, 0xF8, 0xFA, 0xFC, 0xFE}:
        return "documented_prefix_or_extended_form"
    return "core_or_data_candidate"


def scan_asm(path: Path, max_samples: int) -> dict[str, Any]:
    total = 0
    mnems: Counter[str] = Counter()
    unknown_first: Counter[str] = Counter()
    classes: Counter[str] = Counter()
    samples: list[dict[str, str]] = []
    with path.open("r", encoding="utf-8", errors="replace") as fh:
        for line in fh:
            m = INSN_RE.match(line)
            if not m:
                continue
            total += 1
            addr, hex_bytes, mnem = m.groups()
            mnem_l = mnem.lower()
            mnems[mnem_l] += 1
            if not (UNKNOWN_RE.search(line) or UDF_RE.search(line)):
                continue
            first = hex_bytes.split()[0].upper()
            cls = classify_bytes(hex_bytes)
            unknown_first[first] += 1
            classes[cls] += 1
            if len(samples) < max_samples:
                samples.append({"addr": "0x" + addr.upper(), "bytes": " ".join(hex_bytes.split()), "mnemonic": mnem, "class": cls, "line": line.strip()})
    return {"total_insn": total, "mnemonics": dict(mnems.most_common(20)), "unknowns": dict(unknown_first.most_common()), "classes": dict(classes.most_common()), "samples": samples}


def main(argv: list[str]) -> int:
    ap = argparse.ArgumentParser(description="Audit opcode coverage from workspace assembly 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")
    ap.add_argument("--top", type=int, default=0, help="limit number of asm files scanned; 0 means all")
    ap.add_argument("--max-samples", type=int, default=8)
    ap.add_argument("--out", type=Path, default=Path("workspace_opcode_coverage.json"))
    ns = ap.parse_args(argv[1:])
    model_filter = {x.strip().upper() for x in ns.models.split(",")} if ns.models else None
    roots = [x.strip() for x in ns.roots.split(",") if x.strip()]
    asm_files: list[Path] = []
    for root_name in roots:
        root = ns.workspace / root_name
        if root.exists():
            asm_files.extend(sorted(root.rglob("*.asm")))
            asm_files.extend(sorted(root.rglob("*.s")))
    if model_filter:
        asm_files = [p for p in asm_files if infer_model(p).upper() in model_filter]
    if ns.top and ns.top > 0:
        asm_files = asm_files[:ns.top]

    by_model = defaultdict(lambda: {"files": 0, "total_insn": 0, "unknowns": Counter(), "classes": Counter()})
    payloads = []
    for path in asm_files:
        model = infer_model(path)
        scan = scan_asm(path, ns.max_samples)
        d = by_model[model]
        d["files"] += 1
        d["total_insn"] += int(scan["total_insn"])
        d["unknowns"].update(scan["unknowns"])
        d["classes"].update(scan["classes"])
        payloads.append({"model": model, "asm_rel": rel_to(path, ns.workspace), "scan": scan})
    result_models = {}
    for model, d in sorted(by_model.items()):
        classes = dict(d["classes"].most_common())
        actionable = sum(v for k, v in classes.items() if k not in {"break_or_padding", "implemented_fe_bit_operation"})
        result_models[model] = {
            "files": d["files"],
            "total_insn": d["total_insn"],
            "unknowns": dict(d["unknowns"].most_common()),
            "classes": classes,
            "actionable_unknown_count": actionable,
        }
    result = {"schema_version": 2, "source": "workspace_asm_not_profiles", "workspace": str(ns.workspace), "asm_file_count": len(asm_files), "models": result_models, "payloads": payloads}
    ns.out.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    print(f"[OK] audited {len(asm_files)} asm file(s) -> {ns.out}")
    for model, d in result_models.items():
        print(f"[MODEL] {model}: files={d['files']} actionable={d['actionable_unknown_count']} unknown_classes={d['classes']}")
    return 0


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