# ABOUT ME
# Reads the member names of an enum out of a .NET assembly, without needing
# ILSpy, monodis or any other toolchain.
#
# Why it exists: Cosmoteer's DebugMode setting is an enum with no in-game UI and
# no documentation. Guessing a value crashes the game on startup with
# "Deserialization from ... /DebugMode failed" at Enum.Parse, which costs a full
# relaunch per guess. This walks the PE file's CLI metadata tables directly and
# prints the real members, so the value is known before the game is started.
#
# It parses only as far as it must: PE headers -> CLI header -> metadata root ->
# the #~ table stream -> the Module, TypeRef, TypeDef and Field tables.
#
# Usage:  python3 tools/dotnet_enum.py <assembly.dll> [EnumName]
# With no EnumName, lists every type whose name contains "Debug".

import struct
import sys

# Table indices used here.
MODULE, TYPEREF, TYPEDEF, FIELDPTR, FIELD = 0x00, 0x01, 0x02, 0x03, 0x04
METHODDEF = 0x06
MODULEREF, TYPESPEC, ASSEMBLYREF = 0x1A, 0x1B, 0x23


class Reader:
    def __init__(self, data, pos=0):
        self.data = data
        self.pos = pos

    def u16(self):
        value = struct.unpack_from("<H", self.data, self.pos)[0]
        self.pos += 2
        return value

    def u32(self):
        value = struct.unpack_from("<I", self.data, self.pos)[0]
        self.pos += 4
        return value

    def u64(self):
        value = struct.unpack_from("<Q", self.data, self.pos)[0]
        self.pos += 8
        return value


def find_cli_metadata(data):
    """Return (metadata_bytes_offset) for the CLI metadata root."""
    pe_offset = struct.unpack_from("<I", data, 0x3C)[0]
    if data[pe_offset : pe_offset + 4] != b"PE\0\0":
        raise ValueError("not a PE file")

    coff = pe_offset + 4
    number_of_sections = struct.unpack_from("<H", data, coff + 2)[0]
    size_of_optional = struct.unpack_from("<H", data, coff + 16)[0]
    optional = coff + 20
    magic = struct.unpack_from("<H", data, optional)[0]
    # PE32 (0x10b) has a 28-byte windows-specific prefix difference vs PE32+.
    data_dir = optional + (96 if magic == 0x10B else 112)

    sections = []
    section_table = optional + size_of_optional
    for i in range(number_of_sections):
        base = section_table + i * 40
        virtual_size, virtual_address, raw_size, raw_ptr = struct.unpack_from(
            "<IIII", data, base + 8
        )
        sections.append((virtual_address, max(virtual_size, raw_size), raw_ptr))

    def rva_to_offset(rva):
        for virtual_address, size, raw_ptr in sections:
            if virtual_address <= rva < virtual_address + size:
                return raw_ptr + (rva - virtual_address)
        raise ValueError(f"RVA {rva:#x} is not in any section")

    # Data directory entry 14 is the CLI header.
    cli_rva = struct.unpack_from("<I", data, data_dir + 14 * 8)[0]
    cli_offset = rva_to_offset(cli_rva)
    metadata_rva = struct.unpack_from("<I", data, cli_offset + 8)[0]
    return rva_to_offset(metadata_rva)


def parse_streams(data, root):
    if data[root : root + 4] != b"BSJB":
        raise ValueError("CLI metadata signature not found")
    version_length = struct.unpack_from("<I", data, root + 12)[0]
    pos = root + 16 + version_length
    pos += 2  # flags
    stream_count = struct.unpack_from("<H", data, pos)[0]
    pos += 2

    streams = {}
    for _ in range(stream_count):
        offset, size = struct.unpack_from("<II", data, pos)
        pos += 8
        end = data.index(b"\0", pos)
        name = data[pos:end].decode("ascii")
        pos = end + 1
        pos = (pos + 3) & ~3  # names are padded to a 4-byte boundary
        streams[name] = (root + offset, size)
    return streams


class Tables:
    def __init__(self, data, streams):
        self.data = data
        self.strings_base = streams["#Strings"][0]

        base, _size = streams.get("#~") or streams["#-"]
        reader = Reader(data, base)
        reader.u32()  # reserved
        reader.u16()  # major/minor version
        heap_sizes = data[reader.pos]
        reader.pos += 2  # heap sizes + reserved
        valid = reader.u64()
        reader.u64()  # sorted

        self.string_index_size = 4 if heap_sizes & 0x01 else 2
        self.guid_index_size = 4 if heap_sizes & 0x02 else 2
        self.blob_index_size = 4 if heap_sizes & 0x04 else 2

        self.rows = {}
        for table in range(64):
            if valid & (1 << table):
                self.rows[table] = reader.u32()

        self.row_start = reader.pos

    def count(self, table):
        return self.rows.get(table, 0)

    def simple_index(self, table):
        return 4 if self.count(table) >= (1 << 16) else 2

    def coded_index(self, tables, tag_bits):
        largest = max((self.count(t) for t in tables), default=0)
        return 4 if largest >= (1 << (16 - tag_bits)) else 2

    def row_size(self, table):
        s, g, b = self.string_index_size, self.guid_index_size, self.blob_index_size
        if table == MODULE:
            return 2 + s + 3 * g
        if table == TYPEREF:
            return self.coded_index([MODULE, MODULEREF, ASSEMBLYREF, TYPEREF], 2) + 2 * s
        if table == TYPEDEF:
            return (
                4
                + 2 * s
                + self.coded_index([TYPEDEF, TYPEREF, TYPESPEC], 2)
                + self.simple_index(FIELD)
                + self.simple_index(METHODDEF)
            )
        if table == FIELDPTR:
            return self.simple_index(FIELD)
        if table == FIELD:
            return 2 + s + b
        raise ValueError(f"row size for table {table:#x} not implemented")

    def table_offset(self, table):
        offset = self.row_start
        for earlier in sorted(self.rows):
            if earlier == table:
                return offset
            offset += self.row_size(earlier) * self.rows[earlier]
        raise ValueError(f"table {table:#x} not present")

    def string(self, index):
        end = self.data.index(b"\0", self.strings_base + index)
        return self.data[self.strings_base + index : end].decode("utf-8", "replace")

    def read_index(self, reader, size):
        return reader.u32() if size == 4 else reader.u16()

    def typedefs(self):
        """Yield (name, first_field_index) for every TypeDef, 1-based fields."""
        for i, name, first_field, _extends in self.typedefs_full():
            yield i, name, first_field

    def typedefs_full(self):
        """As typedefs(), plus the row index of the base type, or None.

        Extends is a TypeDefOrRef coded index: the low 2 bits are the tag, and
        tag 0 means TypeDef, which is the only case that can be followed within
        this assembly. Base types in other assemblies are reported as None.
        """
        base = self.table_offset(TYPEDEF)
        size = self.row_size(TYPEDEF)
        extends_size = self.coded_index([TYPEDEF, TYPEREF, TYPESPEC], 2)
        field_size = self.simple_index(FIELD)
        for i in range(self.count(TYPEDEF)):
            reader = Reader(self.data, base + i * size)
            reader.u32()  # flags
            name = self.string(self.read_index(reader, self.string_index_size))
            self.read_index(reader, self.string_index_size)  # namespace
            extends = self.read_index(reader, extends_size)
            first_field = self.read_index(reader, field_size)
            base_row = None
            if extends and (extends & 0x3) == 0:
                base_row = (extends >> 2) - 1
            yield i, name, first_field, base_row

    def field_name(self, one_based_index):
        base = self.table_offset(FIELD)
        size = self.row_size(FIELD)
        reader = Reader(self.data, base + (one_based_index - 1) * size)
        reader.u16()  # flags
        return self.string(self.read_index(reader, self.string_index_size))


def clean(field_name):
    """Turn a compiler backing-field name into the property name."""
    if field_name.startswith("<") and field_name.endswith(">k__BackingField"):
        return field_name[1:-len(">k__BackingField")]
    return field_name


def load(assembly_path):
    with open(assembly_path, "rb") as handle:
        data = handle.read()
    return Tables(data, parse_streams(data, find_cli_metadata(data)))


def schema(tables, type_name, follow_bases=True):
    """Return the set of field names declared by a type, including inherited."""
    rows = list(tables.typedefs_full())
    by_name = {name: position for position, (_i, name, _f, _b) in enumerate(rows)}
    field_count = tables.count(FIELD)

    def declared(position):
        _i, _name, first_field, _base = rows[position]
        last_field = rows[position + 1][2] - 1 if position + 1 < len(rows) else field_count
        return {
            clean(tables.field_name(index))
            for index in range(first_field, last_field + 1)
            if index >= 1
        }

    position = by_name.get(type_name)
    if position is None:
        return None

    names = set()
    seen = set()
    while position is not None and position not in seen:
        seen.add(position)
        names |= declared(position)
        position = rows[position][3] if follow_bases else None
    names.discard("value__")
    return names


def main():
    if len(sys.argv) < 2:
        print(__doc__)
        return 2

    with open(sys.argv[1], "rb") as handle:
        data = handle.read()

    tables = Tables(data, parse_streams(data, find_cli_metadata(data)))
    typedefs = list(tables.typedefs())
    field_count = tables.count(FIELD)

    wanted = sys.argv[2] if len(sys.argv) > 2 else None

    for position, (_index, name, first_field) in enumerate(typedefs):
        if wanted is None:
            if "Debug" in name:
                print(name)
            continue
        if name != wanted:
            continue

        # A type's fields run until the next type's first field.
        if position + 1 < len(typedefs):
            last_field = typedefs[position + 1][2] - 1
        else:
            last_field = field_count

        print(f"{name}:")
        for field_index in range(first_field, last_field + 1):
            field = tables.field_name(field_index)
            if field == "value__":
                continue  # the enum's synthetic backing field
            print(f"  {field}")
    return 0


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