"""Minimal PE image and COFF object readers (i386 only, no dependencies)."""
import struct

# i386 COFF relocation types
REL_I386_DIR32 = 0x06
REL_I386_DIR32NB = 0x07
REL_I386_SECTION = 0x0A
REL_I386_SECREL = 0x0B
REL_I386_REL32 = 0x14

IMAGE_SCN_CNT_CODE = 0x20
IMAGE_SCN_LNK_COMDAT = 0x1000


class PEImage:
    def __init__(self, path):
        with open(path, "rb") as f:
            self.data = f.read()
        d = self.data
        pe = struct.unpack_from("<I", d, 0x3C)[0]
        if d[pe:pe + 4] != b"PE\0\0":
            raise ValueError(f"{path}: not a PE file")
        nsec, = struct.unpack_from("<H", d, pe + 6)
        optsize, = struct.unpack_from("<H", d, pe + 20)
        opt = pe + 24
        self.image_base, = struct.unpack_from("<I", d, opt + 28)
        self.sections = []
        for i in range(nsec):
            o = opt + optsize + i * 40
            name = d[o:o + 8].rstrip(b"\0").decode("latin-1")
            vsize, va, rawsize, rawptr = struct.unpack_from("<IIII", d, o + 8)
            chars, = struct.unpack_from("<I", d, o + 36)
            self.sections.append((name, self.image_base + va, vsize, rawptr, rawsize, chars))

    def section(self, name):
        for s in self.sections:
            if s[0] == name:
                return s
        raise KeyError(name)

    def read(self, va, size):
        for name, sva, vsize, rawptr, rawsize, _ in self.sections:
            if sva <= va < sva + max(vsize, rawsize):
                off = va - sva
                out = self.data[rawptr + off:rawptr + min(off + size, rawsize)]
                return out + b"\0" * (size - len(out))  # bss tail
        raise ValueError(f"VA 0x{va:08x} not in any section")

    def contains(self, va):
        return any(s[1] <= va < s[1] + s[2] for s in self.sections)


class CoffSymbol:
    __slots__ = ("index", "name", "value", "section", "type", "storage")

    def __init__(self, index, name, value, section, type_, storage):
        self.index, self.name, self.value = index, name, value
        self.section, self.type, self.storage = section, type_, storage

    @property
    def is_function(self):
        return (self.type >> 4) == 2

    @property
    def is_external(self):
        return self.storage == 2


class CoffSection:
    __slots__ = ("index", "name", "data", "chars", "relocs")

    def __init__(self, index, name, data, chars, relocs):
        self.index, self.name, self.data, self.chars, self.relocs = index, name, data, chars, relocs


IMAGE_COMDAT_SELECT_ANY = 2


class CoffObject:
    def __init__(self, path):
        with open(path, "rb") as f:
            d = f.read()
        self.path = path
        machine, nsec, _, symptr, nsym, optsize, _ = struct.unpack_from("<HHIIIHH", d, 0)
        if machine != 0x14C:
            raise ValueError(f"{path}: not an i386 COFF object")
        strtab = symptr + nsym * 18

        def cstr(raw):
            if raw[:4] == b"\0\0\0\0":
                off, = struct.unpack_from("<I", raw, 4)
                end = d.index(b"\0", strtab + off)
                return d[strtab + off:end].decode("latin-1")
            return raw.rstrip(b"\0").decode("latin-1")

        self.symbols = {}
        self.comdat_select = {}  # section number -> IMAGE_COMDAT_SELECT_*
        i = 0
        while i < nsym:
            o = symptr + i * 18
            raw = d[o:o + 8]
            value, secnum, type_, storage, naux = struct.unpack_from("<IhHBB", d, o + 8)
            self.symbols[i] = CoffSymbol(i, cstr(raw), value, secnum, type_, storage)
            if storage == 3 and value == 0 and naux and secnum > 0:  # section definition
                self.comdat_select[secnum] = d[o + 18 + 14]
            i += 1 + naux

        self.sections = {}
        for s in range(nsec):
            o = 20 + optsize + s * 40
            raw = d[o:o + 8]
            name = raw.rstrip(b"\0").decode("latin-1")
            if name.startswith("/"):
                off = int(name[1:])
                name = d[strtab + off:d.index(b"\0", strtab + off)].decode("latin-1")
            size, rawptr, relptr, _, nrel, _, chars = struct.unpack_from("<IIIIHHI", d, o + 16)
            data = d[rawptr:rawptr + size] if rawptr else b"\0" * size
            relocs = []
            for r in range(nrel):
                va, symidx, typ = struct.unpack_from("<IIH", d, relptr + r * 10)
                relocs.append((va, symidx, typ))
            self.sections[s + 1] = CoffSection(s + 1, name, data, chars, relocs)

    def is_inline(self, secnum):
        """Inline/template functions are emitted as pick-any COMDATs."""
        return self.comdat_select.get(secnum) == IMAGE_COMDAT_SELECT_ANY

    def functions(self):
        """Yield (symbol, section, start, end) for each function defined in a code section."""
        by_sec = {}
        for sym in self.symbols.values():
            if sym.section > 0 and sym.is_function:
                by_sec.setdefault(sym.section, []).append(sym)
        for secnum, syms in by_sec.items():
            sec = self.sections[secnum]
            if not sec.chars & IMAGE_SCN_CNT_CODE:
                continue
            syms.sort(key=lambda s: s.value)
            for i, sym in enumerate(syms):
                end = syms[i + 1].value if i + 1 < len(syms) else len(sec.data)
                yield sym, sec, sym.value, end