"""Dump ACPI firmware tables using GetSystemFirmwareTable.

User-mode, no driver, no admin required.

Subcommands:
    list    enumerate ACPI tables provided by firmware
    dump    save all (or one) ACPI tables to .aml files
    grep    scan saved tables for an ASCII pattern (e.g. MIFS)
"""

from __future__ import annotations

import argparse
import ctypes
import sys
from ctypes import wintypes
from datetime import datetime
from pathlib import Path


def _sig_to_dword(s: str) -> int:
    """4-char ASCII -> DWORD where in-memory bytes spell s (little-endian on x86)."""
    if len(s) != 4:
        raise ValueError("expected exactly 4 ASCII chars")
    return int.from_bytes(s.encode("ascii"), "little")


def _dword_to_sig(value: int) -> str:
    return value.to_bytes(4, "little").decode("ascii", "replace")


def _fourcc_str(value: int) -> str:
    return _dword_to_sig(value)


KERNEL32 = ctypes.WinDLL("kernel32", use_last_error=True)

KERNEL32.EnumSystemFirmwareTables.argtypes = [
    wintypes.DWORD,
    ctypes.c_void_p,
    wintypes.DWORD,
]
KERNEL32.EnumSystemFirmwareTables.restype = wintypes.UINT

KERNEL32.GetSystemFirmwareTable.argtypes = [
    wintypes.DWORD,
    wintypes.DWORD,
    ctypes.c_void_p,
    wintypes.DWORD,
]
KERNEL32.GetSystemFirmwareTable.restype = wintypes.UINT


def _resolve_acpi_provider() -> int:
    candidates = [
        _sig_to_dword("ACPI"),
        int.from_bytes(b"ACPI", "big"),
    ]
    for sig in candidates:
        n = KERNEL32.EnumSystemFirmwareTables(sig, None, 0)
        if n > 0:
            return sig
    raise OSError("EnumSystemFirmwareTables: no ACPI tables reported")


def enum_acpi_tables() -> tuple[int, list[int]]:
    provider = _resolve_acpi_provider()
    needed = KERNEL32.EnumSystemFirmwareTables(provider, None, 0)
    buf = (ctypes.c_uint32 * (needed // 4))()
    written = KERNEL32.EnumSystemFirmwareTables(provider, buf, needed)
    if written == 0:
        raise OSError(
            f"EnumSystemFirmwareTables read failed; GLE={ctypes.get_last_error()}"
        )
    return provider, list(buf[: written // 4])


def get_table(provider: int, table_id: int) -> bytes:
    size = KERNEL32.GetSystemFirmwareTable(provider, table_id, None, 0)
    if size == 0:
        raise OSError(
            f"GetSystemFirmwareTable size query failed for {_fourcc_str(table_id)!r}; "
            f"GLE={ctypes.get_last_error()}"
        )
    buf = (ctypes.c_uint8 * size)()
    written = KERNEL32.GetSystemFirmwareTable(provider, table_id, buf, size)
    if written == 0:
        raise OSError(
            f"GetSystemFirmwareTable read failed for {_fourcc_str(table_id)!r}; "
            f"GLE={ctypes.get_last_error()}"
        )
    return bytes(buf[:written])


def cmd_list(args: argparse.Namespace) -> int:
    provider, ids = enum_acpi_tables()
    print(f"provider=0x{provider:08X} ('ACPI'); {len(ids)} enumerated tables")

    seen_sigs: set[str] = set()
    for index, table_id in enumerate(ids):
        size = KERNEL32.GetSystemFirmwareTable(provider, table_id, None, 0)
        sig = _fourcc_str(table_id)
        seen_sigs.add(sig)
        print(f"  [{index:2d}] sig={sig!r:6} id=0x{table_id:08X} size={size}")

    extras = []
    for extra in ("DSDT", "FACS"):
        if extra not in seen_sigs:
            tid = _sig_to_dword(extra)
            size = KERNEL32.GetSystemFirmwareTable(provider, tid, None, 0)
            if size:
                extras.append((extra, tid, size))
    if extras:
        print("\nadditional fetchable tables not in enum:")
        for sig, tid, size in extras:
            print(f"  sig={sig!r} id=0x{tid:08X} size={size}")
    return 0


def _slugify_signature(table_id: int) -> str:
    raw = table_id.to_bytes(4, "little")
    return "".join(chr(b) if 0x20 <= b < 0x7F else f"_{b:02X}" for b in raw)


def cmd_dump(args: argparse.Namespace) -> int:
    provider, ids = enum_acpi_tables()
    out_dir = Path(args.out_dir).expanduser().resolve()
    out_dir.mkdir(parents=True, exist_ok=True)

    target = args.signature.upper().ljust(4)[:4] if args.signature else None

    enum_set = {_fourcc_str(t) for t in ids}
    extras = []
    for extra in ("DSDT", "FACS"):
        if extra not in enum_set:
            tid = _sig_to_dword(extra)
            if KERNEL32.GetSystemFirmwareTable(provider, tid, None, 0):
                extras.append(tid)
    if target and target not in enum_set and target not in {_fourcc_str(t) for t in extras}:
        tid = _sig_to_dword(target)
        if KERNEL32.GetSystemFirmwareTable(provider, tid, None, 0):
            extras.append(tid)
    ids = list(ids) + extras

    counters: dict[str, int] = {}
    written_count = 0

    for table_id in ids:
        sig_str = _fourcc_str(table_id)
        sig_norm = _slugify_signature(table_id)
        if target and sig_str != target:
            continue
        idx = counters.get(sig_norm, 0)
        counters[sig_norm] = idx + 1

        try:
            data = get_table(provider, table_id)
        except OSError as exc:
            print(f"skip {sig_str!r}: {exc}")
            continue

        suffix = "" if idx == 0 else f"_{idx}"
        path = out_dir / f"{sig_norm}{suffix}.aml"
        path.write_bytes(data)
        oem = data[10:16].decode("ascii", "replace") if len(data) >= 16 else ""
        oem_table = data[16:24].decode("ascii", "replace") if len(data) >= 24 else ""
        print(f"wrote {path.name} ({len(data)} bytes) OEM={oem!r} OEMTable={oem_table!r}")
        written_count += 1

    print(f"\nwrote {written_count} tables to: {out_dir}")
    return 0


def cmd_grep(args: argparse.Namespace) -> int:
    pattern = args.pattern.encode("ascii")
    in_dir = Path(args.in_dir).expanduser().resolve()
    files = sorted(in_dir.glob("*.aml"))
    if not files:
        raise SystemExit(f"no .aml files in {in_dir}")

    total_hits = 0
    for path in files:
        data = path.read_bytes()
        offsets: list[int] = []
        start = 0
        while True:
            idx = data.find(pattern, start)
            if idx < 0:
                break
            offsets.append(idx)
            start = idx + 1
        if offsets:
            total_hits += len(offsets)
            print(f"{path.name}: {len(offsets)} hits at offsets " + ", ".join(f"0x{o:X}" for o in offsets[:8]))
            if args.context:
                for o in offsets[:5]:
                    s = max(0, o - 16)
                    e = min(len(data), o + len(pattern) + 16)
                    chunk = data[s:e]
                    hex_part = " ".join(f"{b:02X}" for b in chunk)
                    ascii_part = "".join(chr(b) if 32 <= b < 127 else "." for b in chunk)
                    print(f"   0x{o:08X}: {ascii_part}")
                    print(f"             {hex_part}")
    print(f"\ntotal: {total_hits} hits across {len(files)} files")
    return 0


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = parser.add_subparsers(dest="cmd", required=True)

    sp = sub.add_parser("list", help="enumerate ACPI tables")
    sp.set_defaults(func=cmd_list)

    sp = sub.add_parser("dump", help="save ACPI tables to .aml files")
    default_out = Path(__file__).resolve().parent / f"acpi_dump_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
    sp.add_argument("--out-dir", default=str(default_out), help=f"output directory; default: {default_out.name}")
    sp.add_argument("--signature", help="only dump tables with this 4-char signature, e.g. DSDT or SSDT")
    sp.set_defaults(func=cmd_dump)

    sp = sub.add_parser("grep", help="ASCII pattern scan inside dumped tables")
    sp.add_argument("pattern", help="ASCII pattern, e.g. MIFS or _SB_")
    sp.add_argument("--in-dir", default=".", help="directory with .aml files")
    sp.add_argument("--context", action="store_true", help="show 16-byte context around each hit")
    sp.set_defaults(func=cmd_grep)

    args = parser.parse_args()
    return args.func(args)


if __name__ == "__main__":
    sys.exit(main())
