"""Minimal ELF64/ELF32 reader (stdlib only).

Needed for three things:
  * symbol resolution for the demo uproc firmware (breakpoints by name),
  * program-header extraction so the real GSP ``.fwimage`` can be placed into
    QEMU RAM (and relocated when a probe needs it),
  * entry-point detection for the first-fault firmware probe.

This is not a full ELF library; it parses headers, PT_LOAD segments and
``.symtab`` symbols, which is all green-eye needs.
"""
from __future__ import annotations

import struct
from dataclasses import dataclass, field as _field

from .errors import GreenEyeError

EM_RISCV = 243
PT_LOAD = 1
SHT_SYMTAB = 2
SHT_STRTAB = 3
STT_FUNC = 2
STT_OBJECT = 1
STT_NOTYPE = 0


class ElfError(GreenEyeError):
    """Malformed or unsupported ELF file."""


@dataclass
class Segment:
    p_type: int
    p_flags: int
    p_offset: int
    p_vaddr: int
    p_paddr: int
    p_filesz: int
    p_memsz: int
    p_align: int

    @property
    def is_load(self) -> bool:
        return self.p_type == PT_LOAD

    @property
    def writable(self) -> bool:
        return bool(self.p_flags & 0x2)

    @property
    def executable(self) -> bool:
        return bool(self.p_flags & 0x1)


@dataclass
class ElfSymbol:
    name: str
    value: int
    size: int
    bind: int
    type: int

    @property
    def is_function(self) -> bool:
        return self.type == STT_FUNC

    @property
    def is_object(self) -> bool:
        return self.type == STT_OBJECT


@dataclass
class FirmwareImage:
    path: str
    is64: bool
    e_machine: int
    entry: int
    segments: list[Segment] = _field(default_factory=list)
    symbols: list[ElfSymbol] = _field(default_factory=list)
    _blob: bytes | None = None

    @property
    def is_riscv(self) -> bool:
        return self.e_machine == EM_RISCV

    @property
    def load_segments(self) -> list[Segment]:
        return [s for s in self.segments if s.is_load]

    @property
    def min_vaddr(self) -> int:
        loads = self.load_segments
        return min((s.p_vaddr for s in loads), default=0)

    @property
    def max_vaddr(self) -> int:
        loads = self.load_segments
        return max((s.p_vaddr + s.p_memsz for s in loads), default=0)

    def read(self) -> bytes:
        if self._blob is None:
            with open(self.path, "rb") as fh:
                self._blob = fh.read()
        return self._blob

    def segment_bytes(self, seg: Segment) -> bytes:
        return self.read()[seg.p_offset : seg.p_offset + seg.p_filesz]

    def flat_load_plan(self, bias: int = 0) -> list[tuple[int, bytes]]:
        """Return ``(load_addr, data)`` pairs for every PT_LOAD segment.

        ``bias`` is added to each segment's virtual address (use it to place a
        firmware linked for GPU physical addresses into QEMU RAM).
        """
        plan: list[tuple[int, bytes]] = []
        for seg in self.load_segments:
            data = bytearray(self.segment_bytes(seg))
            if len(data) < seg.p_memsz:
                data += b"\x00" * (seg.p_memsz - len(data))
            plan.append((seg.p_vaddr + bias, bytes(data)))
        return plan

    def relocated_entry(self, bias: int = 0) -> int:
        return self.entry + bias

    def summary(self) -> str:
        return (
            f"{self.path}: ELF{'64' if self.is64 else '32'} "
            f"machine={self.e_machine} entry={self.entry:#x} "
            f"segments={len(self.segments)} loads={len(self.load_segments)} "
            f"symbols={len(self.symbols)} "
            f"vaddr=[{self.min_vaddr:#x},{self.max_vaddr:#x})"
        )


def load_elf(path: str) -> FirmwareImage:
    with open(path, "rb") as fh:
        data = fh.read()
    return parse_elf(data, path=path)


def parse_elf(data: bytes, path: str = "<memory>") -> FirmwareImage:
    if len(data) < 52 or data[:4] != b"\x7fELF":
        raise ElfError(f"{path}: not an ELF file")
    ei_class = data[4]
    if ei_class == 2:
        is64 = True
        hdr_fmt = "<HHIQQQIHHHHHH"
    elif ei_class == 1:
        is64 = False
        hdr_fmt = "<HHIIIIIHHHHHH"
    else:
        raise ElfError(f"{path}: unknown EI_CLASS {ei_class}")
    hdr_size = struct.calcsize(hdr_fmt)
    fields = struct.unpack_from(hdr_fmt, data, 16)
    (
        _e_type,
        e_machine,
        _e_version,
        e_entry,
        e_phoff,
        e_shoff,
        _e_flags,
        _e_ehsize,
        e_phentsize,
        e_phnum,
        e_shentsize,
        e_shnum,
        e_shstrndx,
    ) = fields

    image = FirmwareImage(path=path, is64=is64, e_machine=e_machine, entry=e_entry)

    if is64:
        ph_fmt = "<IIQQQQQQ"
        sh_fmt = "<IIQQQQIIQQ"
        sym_fmt = "<IBBHQQ"
    else:
        ph_fmt = "<IIIIIIII"
        sh_fmt = "<IIIIIIIIII"
        sym_fmt = "<IIIBBH"

    ph_size = struct.calcsize(ph_fmt)
    for i in range(e_phnum):
        base = e_phoff + i * e_phentsize
        if base + ph_size > len(data):
            break
        p = struct.unpack_from(ph_fmt, data, base)
        image.segments.append(Segment(*p[:8]))

    # Sections -> symbols
    sh_size = struct.calcsize(sh_fmt)
    sections: list[tuple[int, int, int, int]] = []  # type, offset, size, link
    for i in range(e_shnum):
        base = e_shoff + i * e_shentsize
        if base + sh_size > len(data):
            break
        sh = struct.unpack_from(sh_fmt, data, base)
        sh_type, sh_offset, sh_size_val, sh_link = sh[1], sh[4], sh[5], sh[6]
        sections.append((sh_type, sh_offset, sh_size_val, sh_link))

    for sh_type, sh_offset, sh_size_val, sh_link in sections:
        if sh_type != SHT_SYMTAB:
            continue
        strtab_off, strtab_size = 0, 0
        if 0 <= sh_link < len(sections):
            _, strtab_off, strtab_size, _ = sections[sh_link]
        strtab = data[strtab_off : strtab_off + strtab_size]
        count = sh_size_val // struct.calcsize(sym_fmt)
        for i in range(count):
            base = sh_offset + i * struct.calcsize(sym_fmt)
            if base + struct.calcsize(sym_fmt) > len(data):
                break
            st = struct.unpack_from(sym_fmt, data, base)
            if is64:
                st_name, st_info, _st_other, _st_shndx, st_value, st_size_val = st
            else:
                st_name, st_value, st_size_val, st_info, _st_other, _st_shndx = st
            if st_name == 0 or st_value == 0:
                continue
            name = _read_cstr(strtab, st_name)
            if not name:
                continue
            image.symbols.append(
                ElfSymbol(
                    name=name,
                    value=st_value,
                    size=st_size_val,
                    bind=(st_info >> 4) & 0xF,
                    type=st_info & 0xF,
                )
            )
    return image


def _read_cstr(blob: bytes, offset: int) -> str:
    if offset >= len(blob):
        return ""
    end = blob.find(b"\x00", offset)
    if end == -1:
        end = len(blob)
    return blob[offset:end].decode("utf-8", "replace")


def write_relocated_elf(path: str, out_path: str, bias: int) -> str:
    """Write a copy of an ELF64 whose PT_LOADs and entry are shifted by bias.

    ``p_paddr`` is forced to ``p_vaddr + bias`` so the layout matches the
    virtual addresses.  Useful for formats/tools that honour e_entry.
    """
    with open(path, "rb") as fh:
        data = bytearray(fh.read())
    if len(data) < 64 or data[:4] != b"\x7fELF" or data[4] != 2:
        raise ElfError(f"{path}: relocation requires a 64-bit ELF")
    e_entry = struct.unpack_from("<Q", data, 24)[0]
    e_phoff = struct.unpack_from("<Q", data, 32)[0]
    e_phentsize = struct.unpack_from("<H", data, 54)[0]
    e_phnum = struct.unpack_from("<H", data, 56)[0]
    struct.pack_into("<Q", data, 24, e_entry + bias)
    for i in range(e_phnum):
        base = e_phoff + i * e_phentsize
        if base + 56 > len(data):
            break
        p_type = struct.unpack_from("<I", data, base)[0]
        if p_type != PT_LOAD:
            continue
        p_vaddr = struct.unpack_from("<Q", data, base + 16)[0]
        struct.pack_into("<Q", data, base + 16, p_vaddr + bias)
        struct.pack_into("<Q", data, base + 24, p_vaddr + bias)
    with open(out_path, "wb") as fh:
        fh.write(data)
    return out_path


def _trampoline(entry_offset: int) -> bytes:
    """RV64 code at offset 0 that jumps to ``PC + entry_offset``.

    QEMU's RISC-V ``-kernel`` raw/ELF path starts the CPU at the RAM base for
    images whose e_entry it does not honour; this 8-byte stub (auipc t0, hi;
    jalr x0, t0, lo) makes the firmware entry reachable from the base without
    patching the image itself.
    """
    if not -0x80000000 <= entry_offset < 0x80000000:
        raise ElfError("entry offset out of trampoline range")
    hi20 = ((entry_offset + 0x800) >> 12) & 0xFFFFF
    lo12 = entry_offset & 0xFFF
    auipc = (hi20 << 12) | (5 << 7) | 0x17  # auipc t0, hi20
    jalr = (lo12 << 20) | (5 << 15) | 0x67  # jalr x0, t0, lo12
    return struct.pack("<II", auipc, jalr)


def write_flat_kernel(
    path: str,
    out_path: str,
    *,
    load_addr: int = 0x80000000,
    trampoline: bool = True,
) -> tuple[str, int]:
    """Flatten an ELF's PT_LOADs into a raw image QEMU can boot as `-kernel`.

    Returns ``(out_path, entry_addr)`` where ``entry_addr = load_addr +
    (e_entry - min_vaddr)``.  When the entry is not at offset 0 a tiny
    PC-relative trampoline is written at offset 0 (QEMU starts at the load
    base for raw kernels).  ``load_addr`` is not encoded anywhere except in
    the returned entry address: all code is position-independent relative to
    the image base, exactly as it executes on the device.
    """
    image = load_elf(path)
    span = image.max_vaddr - image.min_vaddr
    flat = bytearray(span)
    for seg in image.load_segments:
        start = seg.p_vaddr - image.min_vaddr
        end = start + seg.p_filesz
        flat[start:end] = image.segment_bytes(seg)
    entry_offset = image.entry - image.min_vaddr
    if trampoline and entry_offset != 0:
        stub = _trampoline(entry_offset)
        flat[0 : len(stub)] = stub
    with open(out_path, "wb") as fh:
        fh.write(flat)
    return out_path, load_addr + entry_offset