#!/usr/bin/env python3
"""Extract exact NE segments, entry points, modules, and relocation sites."""

from __future__ import annotations

import argparse
import csv
import struct
from dataclasses import dataclass
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
SOURCE_TYPE_NAMES = {
    0: "low-byte",
    2: "selector-16",
    3: "far-pointer-32",
    5: "offset-16",
    11: "offset-32",
    13: "pointer-48",
}
TARGET_TYPE_NAMES = {
    0: "internal",
    1: "import-ordinal",
    2: "import-name",
    3: "os-fixup",
}


@dataclass(frozen=True)
class Segment:
    number: int
    file_offset: int
    file_length: int
    flags: int
    allocation_size: int

    @property
    def selector(self) -> int:
        return 0x1000 + (self.number - 1) * 8

    def address(self, offset: int) -> str:
        return f"{self.selector:04x}:{offset:04x}"


def u16(data: bytes, offset: int) -> int:
    return struct.unpack_from("<H", data, offset)[0]


def pascal_string(data: bytes, offset: int) -> str:
    length = data[offset]
    return data[offset + 1 : offset + 1 + length].decode("latin-1")


def parse_segments(data: bytes, ne_offset: int) -> list[Segment]:
    count = u16(data, ne_offset + 0x1C)
    table = ne_offset + u16(data, ne_offset + 0x22)
    shift = u16(data, ne_offset + 0x32)
    segments = []
    for index in range(count):
        sector, length, flags, allocation = struct.unpack_from(
            "<HHHH", data, table + index * 8
        )
        segments.append(
            Segment(
                number=index + 1,
                file_offset=sector << shift,
                file_length=length or 0x10000,
                flags=flags,
                allocation_size=allocation or 0x10000,
            )
        )
    return segments


def parse_entries(data: bytes, ne_offset: int) -> dict[int, tuple[int, int, int, str]]:
    table = ne_offset + u16(data, ne_offset + 0x04)
    end = table + u16(data, ne_offset + 0x06)
    ordinal = 1
    entries = {}
    cursor = table
    while cursor < end:
        count = data[cursor]
        segment_indicator = data[cursor + 1]
        cursor += 2
        if count == 0:
            break
        if segment_indicator == 0:
            ordinal += count
            continue
        for _ in range(count):
            flags = data[cursor]
            if segment_indicator == 0xFF:
                int3f = u16(data, cursor + 1)
                segment = data[cursor + 3]
                offset = u16(data, cursor + 4)
                kind = f"movable-int3f-{int3f:04x}"
                cursor += 6
            else:
                segment = segment_indicator
                offset = u16(data, cursor + 1)
                kind = "fixed"
                cursor += 3
            entries[ordinal] = (segment, offset, flags, kind)
            ordinal += 1
    return entries


def parse_modules(data: bytes, ne_offset: int) -> list[str]:
    count = u16(data, ne_offset + 0x1E)
    module_table = ne_offset + u16(data, ne_offset + 0x28)
    import_table = ne_offset + u16(data, ne_offset + 0x2A)
    return [
        pascal_string(data, import_table + u16(data, module_table + index * 2))
        for index in range(count)
    ]


def relocation_sites(
    data: bytes, segment: Segment, first: int, additive: bool
) -> list[int]:
    if additive:
        return [first]
    sites = []
    current = first
    seen = set()
    while current != 0xFFFF:
        if current in seen:
            raise ValueError(
                f"segment {segment.number}: cyclic relocation chain at {current:04x}"
            )
        if current + 2 > segment.file_length:
            raise ValueError(
                f"segment {segment.number}: relocation source outside file image: {current:04x}"
            )
        seen.add(current)
        sites.append(current)
        current = u16(data, segment.file_offset + current)
    return sites


def write_tsv(
    path: Path, columns: tuple[str, ...], rows: list[dict[str, object]]
) -> None:
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(
            handle, fieldnames=columns, delimiter="\t", lineterminator="\n"
        )
        writer.writeheader()
        writer.writerows(rows)


def extract(source: Path, output_dir: Path) -> None:
    data = source.read_bytes()
    ne_offset = struct.unpack_from("<I", data, 0x3C)[0]
    if data[ne_offset : ne_offset + 2] != b"NE":
        raise ValueError(f"not an NE executable: {source}")
    segments = parse_segments(data, ne_offset)
    entries = parse_entries(data, ne_offset)
    modules = parse_modules(data, ne_offset)
    import_table = ne_offset + u16(data, ne_offset + 0x2A)

    segment_rows = []
    for segment in segments:
        segment_rows.append(
            {
                "segment": segment.number,
                "selector": f"{segment.selector:04x}",
                "file_offset": segment.file_offset,
                "file_length": segment.file_length,
                "allocation_size": segment.allocation_size,
                "zero_fill_bytes": segment.allocation_size - segment.file_length,
                "flags": f"0x{segment.flags:04x}",
                "has_relocations": bool(segment.flags & 0x0100),
            }
        )

    entry_rows = []
    for ordinal, (segment, offset, flags, kind) in sorted(entries.items()):
        selector = 0x1000 + (segment - 1) * 8
        entry_rows.append(
            {
                "ordinal": ordinal,
                "segment": segment,
                "offset": f"0x{offset:04x}",
                "address": f"{selector:04x}:{offset:04x}",
                "flags": f"0x{flags:02x}",
                "kind": kind,
            }
        )

    module_rows = [
        {"module_index": index, "module_name": name}
        for index, name in enumerate(modules, 1)
    ]

    relocation_rows = []
    for segment in segments:
        if not (segment.flags & 0x0100):
            continue
        cursor = segment.file_offset + segment.file_length
        count = u16(data, cursor)
        cursor += 2
        for record_index in range(1, count + 1):
            source_type, flags, first_source, target1, target2 = struct.unpack_from(
                "<BBHHH", data, cursor
            )
            cursor += 8
            target_type = flags & 0x03
            additive = bool(flags & 0x04)
            target_kind = TARGET_TYPE_NAMES.get(target_type, f"target-{target_type}")
            target = ""
            target_address = ""
            module = ""
            import_ordinal = ""
            import_name = ""
            if target_type == 0:
                if target1 == 0x00FF:
                    entry = entries.get(target2)
                    target = f"entry-ordinal:{target2}"
                    if entry:
                        target_address = segments[entry[0] - 1].address(entry[1])
                else:
                    target = f"segment:{target1}:offset:{target2:04x}"
                    target_address = segments[target1 - 1].address(target2)
            elif target_type == 1:
                module = modules[target1 - 1]
                import_ordinal = target2
                target = f"{module}.ordinal:{target2}"
            elif target_type == 2:
                module = modules[target1 - 1]
                import_name = pascal_string(data, import_table + target2)
                target = f"{module}.{import_name}"
            else:
                target = f"os-fixup:{target1}:{target2}"

            sites = relocation_sites(data, segment, first_source, additive)
            for chain_index, site in enumerate(sites):
                relocation_rows.append(
                    {
                        "segment": segment.number,
                        "record": record_index,
                        "chain_index": chain_index,
                        "source_offset": f"0x{site:04x}",
                        "source_address": segment.address(site),
                        "source_type": SOURCE_TYPE_NAMES.get(
                            source_type & 0x0F, f"source-{source_type & 0x0F}"
                        ),
                        "additive": additive,
                        "target_kind": target_kind,
                        "target": target,
                        "target_address": target_address,
                        "module": module,
                        "import_ordinal": import_ordinal,
                        "import_name": import_name,
                        "raw_flags": f"0x{flags:02x}",
                    }
                )

    write_tsv(
        output_dir / "NE_SEGMENTS.tsv",
        (
            "segment",
            "selector",
            "file_offset",
            "file_length",
            "allocation_size",
            "zero_fill_bytes",
            "flags",
            "has_relocations",
        ),
        segment_rows,
    )
    write_tsv(
        output_dir / "ENTRY_POINTS.tsv",
        ("ordinal", "segment", "offset", "address", "flags", "kind"),
        entry_rows,
    )
    write_tsv(
        output_dir / "MODULE_REFERENCES.tsv",
        ("module_index", "module_name"),
        module_rows,
    )
    write_tsv(
        output_dir / "RELOCATIONS.tsv",
        (
            "segment",
            "record",
            "chain_index",
            "source_offset",
            "source_address",
            "source_type",
            "additive",
            "target_kind",
            "target",
            "target_address",
            "module",
            "import_ordinal",
            "import_name",
            "raw_flags",
        ),
        relocation_rows,
    )
    print(f"segments={len(segment_rows)}")
    print(f"entry_points={len(entry_rows)}")
    print(f"module_references={len(module_rows)}")
    print(f"relocation_sites={len(relocation_rows)}")


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("source", nargs="?", type=Path, default=ROOT / "TDKPIN.EXE")
    parser.add_argument("--output-dir", type=Path, default=ROOT)
    args = parser.parse_args()
    extract(args.source.resolve(), args.output_dir.resolve())
    return 0


if __name__ == "__main__":
    raise SystemExit(main())