import re
import sys
from pathlib import Path

def parse_values(ntstatus_h: Path) -> dict[int,str]:
    if not ntstatus_h.is_file():
        print(f'ERROR: {ntstatus_h} is not a file')
        sys.exit(2)

    values: dict[int,str] = dict()
    lines = ntstatus_h.read_text().splitlines()
    for line in lines:
        # Typical NTSTATUS definition lines look like this:
        # #define STATUS_PENDING                   ((NTSTATUS)0x00000103L)    // winnt
        # The regex will miss a couple STATUS_* values by design:
        # - STATUS_SEVERITY_* (they are not real NTSTATUS values)
        # - STATUS_FLT_DISALLOW_FSFILTER_IO (which is an alias for STATUS_FLT_DISALLOW_FAST_IO)
        m = re.fullmatch(r'\s*#define\s+(STATUS_\w+)\s+\(\(NTSTATUS\)(0x[0-9A-F]+)L\)\s*(//.*)?', line)
        if not m:
            continue

        name = m.group(1)
        value = int(m.group(2), 16)

        # We should remove STATUS_WAIT_num, but not STATUS_WAIT_* (because of STATUS_WAIT_FOR_OPLOCK)
        if re.fullmatch(r'STATUS_WAIT_\d+', name):
            continue

        # We should remove STATUS_ABANDONED_WAIT_num, but not STATUS_ABANDONED_WAIT
        if re.fullmatch(r'STATUS_ABANDONED_WAIT_\d+', name):
            continue

        if value in values:
            print(f'ERROR: duplicated value: {name}=0x{value:08X}, previously seen as {values[value]}')
            sys.exit(2)

        values[value] = name
    return values

def read_script(ntstatus_py: Path) -> tuple[list[str],str]:
    if not ntstatus_py.is_file():
        print(f'ERROR: {ntstatus_py} is not a file')
        sys.exit(2)

    autogen_line = '# AUTOGENERATED FROM HERE, DO NOT TOUCH'
    lines = ntstatus_py.read_text().splitlines()
    for index, line in enumerate(lines):
        if m := re.fullmatch(r'^(\s+)' + autogen_line, line):
            return lines[0:index+1] + [''], m.group(1)

    print(f'ERROR: {ntstatus_py} does not contain the line "{autogen_line}" (with indentation)')
    sys.exit(2)

def main() -> None:
    if len(sys.argv) != 2:
        print(f'Usage: {Path(__file__).name} path/to/ntstatus.h')
        sys.exit(1)

    ntstatus_h = Path(sys.argv[1])
    values = parse_values(ntstatus_h)

    ntstatus_py = Path(__file__).parent.parent / 'ntstatus' / '__init__.py'

    lines, indent = read_script(ntstatus_py)
    for value, name in sorted(values.items()):
        lines.append(f'{indent}{name} = ThirtyTwoBits(0x{value:08X})')

    ntstatus_py.write_text('\n'.join(lines) + '\n')

if __name__ == '__main__':
    main()
