import subprocess
import tempfile
from typing import Literal, IO, Generator
from dataclasses import dataclass, field
from contextlib import contextmanager
import logging
from inspect import cleandoc

from elftools.elf.elffile import ELFFile

from .core import RaspsimException, Prot, getProtFromELFSegment, Raspsim

# "-g",
# "-gdwarf-4",
# "-O3",
# "-Wall",
# "-Wextra",
# "-Wpedantic",
# "-pedantic-errors",

# if segment["p_flags"] & P_FLAGS.PF_W:
#    raise ValueError("Invalid ELF: contains writable segment")
# if elf_data:
#    raise ValueError("Invalid ELF: too many load segments")
# if not segment["p_flags"] & P_FLAGS.PF_X:
#    raise ValueError("Invalid ELF: no executable segment")
# if elf.entry < 0x3000:
#     raise RaspsimException("vaddr must be >= 0x3000")


@contextmanager
def rscompile(
    code: str,
    entry_label: str = "_start",
    lang: Literal["c", "assembler"] = "assembler",
    compiler: str = "cc",
    std: str = "c17",
    text_segment_addr: int = 0x4000,
    additional_flags: list[str] | None = None,
) -> Generator[IO[bytes], None, None]:
    """
    Compile assembly code to an ELF binary.

    Args:
        code: Assembly code to compile.
        entry_label: Entry point label.
        lang: Language of the code.
        compiler: Compiler to use.
        std: Standard to use.
        additional_flags: Additional flags to pass to the compiler.

    Yields:
        A file-like object containing the compiled ELF binary.
    """
    args = [
        compiler,
        f"-x{lang}",
        "-pipe",
        f"-std={std}",
        "-static",
        "-march=k8",
        "-mno-mmx",
        "-mno-80387",
        "-mtune=generic",  # prevent excessive loop unrolling
        "-ffreestanding",
        "-nostdlib",
        "-D_MM_MALLOC_H_INCLUDED",
        # Assembler arguments
        "-Wa,-march=k8+cmov+nommx",
        "-Wa,-mx86-used-note=no",
        "-Wa,--fatal-warnings",
        # Linker arguments
        "-Wl,-n",
        "-Wl,--fatal-warnings",  # fail if entry point doesn't exist
        "-Wl,--no-dynamic-linker",
        "-Wl,--build-id=none",
        "-Wl,-z,defs",
        "-Wl,-z,noexecstack",
        "-Wl,-z,norelro",
        "-Wl,-z,noseparate-code",
        f"-Wl,-Ttext={text_segment_addr:#x}",
        f"-Wl,-e,{entry_label}",
    ] + (additional_flags or [])
    with tempfile.NamedTemporaryFile("w+b") as f:
        subprocess.run(
            args + ["-o", f.name, "-"],
            input=code.encode(),
            check=True,
        )
        f.seek(0)
        yield f


@dataclass
class Segment:
    """
    A class representing a memory segment.
    Attributes:
        vaddr (int): The virtual address of the segment.
        prot (Prot): The protection attributes of the segment.
        size (int): The size of the segment.
        data (bytes | None): Optional data contained in the segment.
    Methods:
        __repr__(): Returns a string representation of the Segment instance.
    Note:
        Length of `data` shall not exceed the size of the segment.
        If size is greater than the length of the data, the rest of the segment denotes zeroed memory.
    """

    vaddr: int
    prot: Prot
    size: int
    data: bytes | None = None

    def __repr__(self):
        data_repr = f" data={self.data!r}" if self.data else ""
        return f"Segment(vaddr={self.vaddr:#x}, prot={self.prot}, size={self.size:#x}{data_repr})"


@dataclass
class ELF:
    """
    ELF binary data.

    Attributes:
        entry (int): The entry point of the ELF binary.
        stack (int | None): The stack pointer, if available.
        segments (list[Segment]): A list of segments loaded into memory.
    Methods:
        __repr__(): Returns a string representation of the ELF object.
    Note:
        Only contains the segments that are loaded into memory. Requires x86_64 executable ELF.
    """

    entry: int
    stack: int | None = None
    segments: list[Segment] = field(default_factory=list)

    def __repr__(self):
        return f"ELF(entry={self.entry:#x}, stack={self.stack:#x}, segments={self.segments})"


class InvalidELFException(Exception):
    """
    Exception raised when an invalid ELF file is encountered.
    """


def load_elf(
    elf: IO[bytes], gnu_stack_size: int = 0x1000, gnu_stack_addr: int = 0x7FE000000000
) -> ELF:
    """
    Load an ELF file and parse its segments.

    Args:
        elf (IO[bytes]): A byte stream representing the ELF file.
        gnu_stack_size (int, optional): The size of the stack when a GNU Stack Segment is found. Defaults to 0x1000.
        gnu_stack_addr (int, optional): The address of the segment when a GNU Stack Segment is found. Defaults to 0x7FE000000000.

    Returns:
        ELF: An ELF object containing the parsed segments and entry point.

    Raises:
        InvalidELFException: If the ELF file is not of type ET_EXEC or machine type EM_X86_64.
        InvalidELFException: If a segment's file size is greater than its memory size.
        ValueError: If the segment data length does not match the segment's file size.
        InvalidELFException: If an unhandled segment type is encountered.
    """
    # new stream from bytes
    elffile = ELFFile(elf)

    if elffile["e_type"] != "ET_EXEC":
        raise InvalidELFException(f"Invalid ELF type {elffile['e_type']}")

    if elffile["e_machine"] != "EM_X86_64":
        raise InvalidELFException(f"Invalid ELF machine {elffile['e_machine']}")

    loaded_elf = ELF(entry=elffile["e_entry"])

    for segment in elffile.iter_segments():
        logging.debug(
            "[.] ELF segment: %s V=%#x FSZ=%#x MSZ=%#x flags=%d",
            segment["p_type"],
            segment["p_vaddr"],
            segment["p_filesz"],
            segment["p_memsz"],
            segment["p_flags"],
        )
        if segment["p_type"] == "PT_LOAD":
            if segment["p_filesz"] > segment["p_memsz"]:
                raise InvalidELFException(
                    "Invalid ELF: corrupt binary. p_filesz > p_memsz"
                )

            # Try to read segment data
            segment_data = segment.data()
            if len(segment_data) != segment["p_filesz"]:
                raise ValueError("Invalid ELF: corrupt binary")

            logging.debug("[.] Read %d bytes of code/data", len(segment_data))
            loaded_elf.segments.append(
                Segment(
                    vaddr=segment["p_vaddr"],
                    size=segment["p_memsz"],
                    data=segment_data,
                    prot=getProtFromELFSegment(segment["p_flags"]),
                )
            )
        elif segment["p_type"] == "PT_GNU_STACK":
            loaded_elf.stack = gnu_stack_addr
            loaded_elf.segments.append(
                Segment(
                    vaddr=gnu_stack_addr - gnu_stack_size,
                    size=gnu_stack_size,
                    prot=Prot.RW,
                )
            )
        elif segment["p_type"] == "PT_NULL":
            pass
        else:
            raise InvalidELFException(
                f"Invalid ELF: unhandled segment {segment['p_type']}"
            )

    return loaded_elf


def elf_add_trampoline(elf: ELF):
    """
    Adds a trampoline to the ELF binary.
    
    The trampoline is a call instruction to the entry point followed by an terminate instruction.
    If reaching the terminator indicates that the program has gracefully returned from the call.
    
    Parameters:
        elf (ELF): The ELF object to which the trampoline will be added.
    Raises:
        RaspsimException: If the entry point is not page-aligned or if the entry 
                          point is less than 0x3000.
    
    Modifies:
        - Adjusts the entry point of the ELF object.
        - Appends a new segment containing the trampoline code to the ELF object's segments.
    """
    if elf.entry & 0xFFF:
        raise RaspsimException("code must be page-aligned")

    if elf.entry < 0x3000:
        raise RaspsimException("vaddr must be >= 0x3000 to use trampoline")

    trampoline_bias = 0x2000
    addr_start = max(0x1000, elf.entry - trampoline_bias)
    reloff = elf.entry - addr_start - 5
    trampoline = b"\xe8" + reloff.to_bytes(4, "little") + b"\xcd\x80"
    elf.segments.append(
        Segment(
            vaddr=addr_start, prot=Prot.RX, size=len(trampoline), data=trampoline
        )
    )
    elf.entry = addr_start

def asm_preamble(label: str = "_start") -> str:
    return cleandoc(f"""
                    .intel_syntax noprefix
                    .global {label}
                    .section .text
                    {label}:
                    """)

def asm_stop_sim() -> str:
    return "int 0x80\n"

def populate_from_elf(sim: Raspsim, elf: ELF, abi: Literal["sysv"] = "sysv"):
    """
    Populate the simulator with the ELF binary data.

    Args:
        sim: Raspsim instance.
        elf: ELF binary data.
        abi: ABI to use. Currently only supports sysv.
    """
    del abi  # unused

    sim.rip = elf.entry

    if elf.stack is None:
        raise RaspsimException("stack must be set")
    sim.rsp = elf.stack

    for segment in elf.segments:
        sim.memmap(
            segment.vaddr, segment.prot, length=segment.size, data=segment.data or b""
        )
