"""

makerel.py - Generate .rel files from .plf files and a static binary

"""

import sys
import glob
import os
import traceback

from pathlib import Path
from dataclasses import dataclass, field
from typing import List, Set, Tuple, Dict

try:
    import libelf
    import librel
    import click
    import logging

    from rich.logging import RichHandler
    from rich.console import Console
except ImportError as e:
    MISSING_PREREQUISITES = (
        f"Missing prerequisite python module {e}.\n"
        f"Run `python3 -m pip install --user -r tools/requirements.txt` to install prerequisites."
    )

    print(MISSING_PREREQUISITES, file=sys.stderr)
    sys.exit(1)

VERSION = "1.0"
CONSOLE = Console()

logging.basicConfig(
    level="NOTSET",
    format="%(message)s",
    datefmt="[%X]",
    handlers=[RichHandler(console=CONSOLE, rich_tracebacks=True)],
)

LOG = logging.getLogger("rich")
LOG.setLevel(logging.INFO)

SECTION_MASK = {".init", ".text", ".ctors", ".dtors", ".rodata", ".data", ".bss"}

REL_SECTION_MASK = {
    ".rela.init",
    ".rela.text",
    ".rela.ctors",
    ".rela.dtors",
    ".rela.rodata",
    ".rela.data",
    ".rela.bss",
}


@click.group()
@click.version_option(VERSION)
def makerel():
    pass


@makerel.command(name="unresolved")
@click.option("--debug/--no-debug")
@click.option("--output", "-o", default="forceactive.txt", required=True)
@click.argument("str_paths", metavar="<ELFs>", nargs=-1)
def unresolved(debug, output, str_paths):
    """Generate a list of symbols which must be in the static executable (and other RELs)."""

    if debug:
        LOG.setLevel(logging.DEBUG)

    static, plfs = load_elfs(str_paths)

    if static:
        LOG.error(f"unresolved does not handle executable files '{static.path}'")
        sys.exit(1)

    undef_symbols = set()
    for plf in plfs:
        for relocation in plf.relocations:
            symbol = relocation.symbol
            if isinstance(symbol, libelf.UndefSymbol):
                assert symbol.name
                undef_symbols.add(symbol.name)

    LOG.debug(f"found {len(undef_symbols)} symbols")

    names = list(undef_symbols)
    names.sort()
    with open(output, "w") as file:
        for name in names:
            file.write(f"{name}\n")

    CONSOLE.LOG(f"Unresolved Symbols: '{output}'")


@makerel.command(name="build")
@click.option("--debug/--no-debug")
@click.option("--yaz0", "-y", "compress_yaz0")
@click.option("--id-offset", "-i", "rel_id_offset", default=1)
@click.option(
    "--spoof-path", "-q", "spoof_path", default="D:\\zeldaGC_USA\\dolzel2\\bin\\Final\\"
)
@click.option("--string-table", "-s", "string_path", required=True)
@click.option(
    "--symbols", default="ELF", type=click.Choice(["ELF", "DEFS"], case_sensitive=False)
)
@click.argument("str_paths", metavar="<ELFs>", nargs=-1)
def build(
    debug, symbols, str_paths, rel_id_offset, compress_yaz0, spoof_path, string_path
):
    """Build RELs files from a list of plfs files."""

    if debug:
        LOG.setLevel(logging.DEBUG)

    static, plfs = load_elfs(str_paths)
    if not static:
        LOG.error(f"static executable ('main.elf') expected")
        sys.exit(1)

    #
    id = rel_id_offset
    elfs = []
    for plf in plfs:
        if static == plf:
            rel = IndexedElf(0, plf)
        else:
            rel = IndexedElf(id, plf)
            id += 1
        elfs.append(rel)

    # sort relocations
    for elf in elfs:
        if "_unresolved" in elf.plf.symbol_map:
            elf._unresolved = elf.plf.symbol_map["_unresolved"][0]

        for _, relocations in elf.plf.section_relocations:
            relocations.sort(key=lambda r: r.offset)

    # symbol table
    symbol_table = dict()
    for elf in elfs:
        for symbol in elf.plf.symbols:
            if isinstance(symbol, libelf.UndefSymbol):
                continue
            if not symbol.isBindGlobal():
                continue
            symbol_table[symbol.name] = (elf, symbol)

    for elf in elfs:
        if elf.id != 0:
            elf.resolve_relocations(symbol_table)

    string_list = StringList(string_path, spoof_path)
    for elf in elfs:
        if elf.id != 0:
            write_rel_from_elf(elf, string_list, compress_yaz0)

    string_list.write()


def apply_rel24_relocation(relocation, section, symbol):
    if not symbol or not isinstance(symbol, libelf.OffsetSymbol):
        return False

    try:
        if librel.apply_relocation(
            relocation.type,
            0,
            section.data,
            0,
            relocation.offset,
            symbol.offset,
            relocation.addend,
        ):
            return True
    except librel.RELRelocationException as e:
        LOG.error(f"applying relocation failed!")
        LOG.error(
            f"relocation: {librel.RELOCATION_NAMES[relocation.type]} {relocation.offset:04X}"
        )
        LOG.error(
            f"section: {section.header.sh_addr:08X} {section.header.sh_size:04X} {section.name}"
        )
        LOG.error(f"symbol: {symbol.offset:08X} {symbol.name}+0x{relocation.addend:X}")
        LOG.error(e)
        CONSOLE.print_exception()
        traceback.print_tb()
        traceback.print_exc()
        sys.exit(1)

    return False


@dataclass
class StringList:
    output_path: str
    spoof_path: str
    data: str = ""

    def add(self, path):
        offset = len(self.data)
        name = str(path)
        if self.spoof_path:
            name = f"{self.spoof_path}{path.name} "
        self.data += name
        return offset, len(name)

    def write(self):
        if len(self.data) > 0:
            with open(self.output_path, "w") as file:
                file.write(self.data)


@dataclass
class ImpTable:
    id: int
    last_section: int
    section_offset: int = 0
    relocations: List[Tuple[int, int, int, int]] = field(default_factory=list)
    rel_offset: int = 0

    def section(self, section_id):
        self.last_section = section_id
        self.section_offset = 0
        self.relocations.append((0, librel.R_DOLPHIN_SECTION, section_id, 0))

    def end(self):
        self.relocations.append((0, librel.R_DOLPHIN_END, 0, 0))

    def nop(self, offset):
        self.relocations.append((offset, librel.R_DOLPHIN_NOP, 0, 0))

    def relocation(self, offset, type, section, addend):
        self.relocations.append((offset, type, section, addend))


@dataclass
class IndexedElf:
    id: int
    plf: libelf.Object

    _unresolved: libelf.Symbol = None
    imp_tables: Dict[int, ImpTable] = field(default_factory=dict)
    imp_table_order: List[int] = field(default_factory=list)
    complete_relocations: Set[libelf.Relocation] = field(default_factory=set)

    def symbol_by_name(self, name):
        if name in self.plf.symbol_map:
            symbols = self.plf.symbol_map[name]
            if len(symbols) > 1:
                LOG.warning(symbols)
            return symbols[0]
        return None

    def resolve_relocation(self, section, relocation, symbol_table):
        # relocations of function calls should all point to _unresolved if it is defined.
        if self._unresolved:
            if relocation.type == librel.R_PPC_REL24:
                apply_rel24_relocation(relocation, relocation.section, self._unresolved)

        ext_symbol = None
        my_symbol = relocation.symbol

        if my_symbol.name and my_symbol.name.endswith("stringBase0"):
            replace_symbol = self.symbol_by_name("@stringBase0")
            if replace_symbol:
                my_symbol = replace_symbol

        # LOG.info(f"relocation for: {my_symbol.name} ({type(my_symbol).__name__})")
        other = None
        if isinstance(my_symbol, libelf.UndefSymbol):
            if my_symbol.name in symbol_table:
                other, ext_symbol = symbol_table[my_symbol.name]
        else:
            other = self
            ext_symbol = my_symbol
            if relocation.type == librel.R_PPC_REL24:
                if apply_rel24_relocation(relocation, relocation.section, ext_symbol):
                    self.complete_relocations.add(relocation)
                    return True

        if ext_symbol:
            # different module id, create new imp table
            if not other.id in self.imp_tables:
                self.imp_table_order.append(other.id)
                self.imp_tables[other.id] = ImpTable(other.id, -1)

            table = self.imp_tables[other.id]

            # new section
            if table.last_section != section.header.sh_info:
                assert section.header.sh_info <= 255
                table.section(section.header.sh_info)

            relative_offset = relocation.offset - table.section_offset
            table.section_offset = relocation.offset

            # too big offset, insert nop's until the offset is manageable
            while relative_offset >= 0xFFFF:
                table.nop(0xFFFF)
                relative_offset -= 0xFFFF

            if not ext_symbol.section:
                assert other.id == 0
                assert isinstance(ext_symbol, libelf.AbsoluteSymbol)

                found_section = None
                for section in other.plf.sections.values():
                    if (
                        ext_symbol.address >= section.header.sh_addr
                        and ext_symbol.address
                        < section.header.sh_addr + section.header.sh_size
                    ):
                        found_section = section
                        break

                if found_section:
                    ext_symbol.section = found_section
                else:
                    LOG.error(
                        f"error no-section provided for relocation of: {my_symbol.name} ({type(my_symbol).__name__}) ({self.plf.name} <- {other.plf.name})"
                    )
                    LOG.error(vars(relocation))
                    LOG.error(ext_symbol)
                    k = vars(ext_symbol)
                    k.pop("object", None)
                    k.pop("section", None)
                    LOG.error(k)
                    sys.exit(1)
            assert ext_symbol.section.header.sh_info <= 255

            addend = relocation.addend
            if isinstance(ext_symbol, libelf.OffsetSymbol):
                addend += ext_symbol.offset
            elif isinstance(ext_symbol, libelf.AbsoluteSymbol):
                addend += ext_symbol.address

            section_id = ext_symbol.section.header.sh_info
            if section_id == 0:
                section_id = ext_symbol.section.header.id

            table.relocation(relative_offset, relocation.type, section_id, addend)
            self.complete_relocations.add(relocation)
            return True
        return False

    def resolve_relocations(self, symbol_table):
        for name, relocations in self.plf.section_relocations:
            if not name in REL_SECTION_MASK:
                continue
            section = self.plf.sections[name]
            for relocation in relocations:
                if not self.resolve_relocation(section, relocation, symbol_table):
                    LOG.error(f"relocation failed: {name:<14} {relocation.symbol.name}")
                    sys.exit(1)


def align_next(offset, alignment):
    return (offset - 1 + alignment) & ~(alignment - 1)


def write_rel(
    path: Path,
    id: int,
    align: int,
    bss_align: int,
    bss_size: int,
    name_offset: int,
    name_size: int,
    prolog: libelf.Symbol,
    epilog: libelf.Symbol,
    unresolved: libelf.Symbol,
    sections: List[librel.Section],
    imp_tables: List[ImpTable],
):

    output = librel.REL()
    output.index = id
    output.numSections = len(sections)
    output.sectionInfoOffset = 0x4C  # for version 3
    output.nameOffset = name_offset
    output.nameSize = name_size
    output.version = 3
    output.bssSize = bss_size
    output.relOffset = 0
    output.impOffset = 0
    output.impSize = 0

    if prolog:
        assert isinstance(prolog, libelf.OffsetSymbol)
        output.prolog = prolog.offset
        output.prologSection = prolog.section.header.id

    if epilog:
        assert isinstance(epilog, libelf.OffsetSymbol)
        output.epilog = epilog.offset
        output.epilogSection = epilog.section.header.id

    if unresolved:
        assert isinstance(unresolved, libelf.OffsetSymbol)
        output.unresolved = unresolved.offset
        output.unresolvedSection = unresolved.section.header.id

    assert output.version >= 2
    output.align = align
    output.bssAlign = bss_align

    assert output.version >= 3

    with path.open("wb") as file:
        librel.write_header(file, output)

        sections_offset = file.tell()
        assert sections_offset == output.sectionInfoOffset

        for section in sections:
            librel.write_section_header(file, section)

        section_data_offset = file.tell()
        for section in sections:
            if section.data:
                padding = section.offset - file.tell()
                if padding > 0:
                    file.write(b"\x00" * padding)

                assert section.offset == file.tell()
                librel.write_section_data(file, section)

        output.impOffset = file.tell()
        output.impSize = len(imp_tables) * 0x8

        file.write(b"\xFF" * output.impSize)

        output.fixSize = file.tell()
        output.relOffset = file.tell()
        rel_offset = output.relOffset
        for table in imp_tables:
            table.rel_offset = rel_offset
            if table.id == id:
                output.fixSize = rel_offset
            for relocation in table.relocations:
                librel.write_relocation(file, *relocation)
            rel_offset += len(table.relocations) * 8

        file.seek(output.impOffset, os.SEEK_SET)
        for table in imp_tables:
            librel.write_imp(file, table.id, table.rel_offset)

        file.seek(0, os.SEEK_SET)
        librel.write_header(file, output)


def write_rel_from_elf(elf: IndexedElf, string_list: StringList, compress_yaz0: bool):
    assert elf.id != 0

    path = Path(str(elf.plf.path).replace(".plf", ".rel"))
    yaz0_path = Path(str(elf.plf.path).replace(".plf", ".rel.yaz0"))
    LOG.debug(f"output: '{path}'")

    # add the plf path to the string list
    name_offset, name_size = string_list.add(elf.plf.path)

    # insert end relocation value for each imp-table
    for table in elf.imp_tables.values():
        if len(table.relocations) > 0:
            table.end()

    # count sections
    section_count = 1  # null section
    for elf_section in elf.plf.sections.values():
        if elf_section.name == ".dead" or elf_section.name == ".rela.dead":
            continue
        section_count += 1

    bss_size = 0
    bss_align = 1
    align = 1
    offset = 0x4C + 0x8 * section_count

    # build section list based on elf sections
    sections = list()
    sections.append(librel.Section(0, 0, False, False, 0))
    for elf_section in elf.plf.sections.values():
        if not elf_section.name in SECTION_MASK:
            continue

        section = librel.Section(
            elf_section.header.id, 0, False, False, elf_section.header.sh_size
        )
        if elf_section.header.sh_type == libelf.SHT_NOBITS:
            if elf_section.header.sh_addralign >= 1:
                if elf_section.header.sh_addralign >= bss_align:
                    bss_align = elf_section.header.sh_addralign

            bss_size += section.length
        else:
            if elf_section.header.sh_addralign >= 1:
                offset = align_next(offset, elf_section.header.sh_addralign)
                assert offset % elf_section.header.sh_addralign == 0

                if elf_section.header.sh_addralign >= align:
                    align = elf_section.header.sh_addralign

            section.offset = offset
            section.data = elf_section.data
            if (elf_section.header.sh_flags & libelf.SHF_EXECINSTR) != 0:
                section.executable_flag = True

            offset += section.length

        sections.append(section)

    while len(sections) < section_count:
        sections.append(librel.Section(len(sections), 0, False, False, 0))

    # find prolog, epilog, and unresolved symbols
    prolog = elf.symbol_by_name("_prolog")
    epilog = elf.symbol_by_name("_epilog")
    unresolved = elf.symbol_by_name("_unresolved")

    # the order of the imp-tables are:
    #       sorted rels (by module id) except the current one and module 0
    #       imp-table for the current rel
    #       imp-table for module 0
    tables = []
    for table in elf.imp_tables.values():
        if table.id == 0 or table.id == elf.id:
            continue
        tables.append(table)
    tables.sort(key=lambda x: x.id)

    if elf.id in elf.imp_tables:
        tables.append(elf.imp_tables[elf.id])

    if 0 in elf.imp_tables:
        tables.append(elf.imp_tables[0])

    # write the rel files
    write_rel(
        path,
        elf.id,
        align,
        bss_align,
        bss_size,
        name_offset,
        name_size,
        prolog,
        epilog,
        unresolved,
        sections,
        tables,
    )


def load_elfs(str_paths):
    static = None
    plfs = []
    for plfs_path in str_paths:
        files = glob.glob(plfs_path, recursive=True)
        if len(files) <= 0:
            LOG.error(f"found zero files for path: '{plfs_path}'")

        for plf_path in files:
            path = Path(plf_path)
            LOG.debug(f"load: '{path}'")
            try:
                plf = libelf.load_object_from_path(path)
                if plf.executable:
                    if static:
                        LOG.error(f"multiple executable files not supported")
                        sys.exit(1)
                    static = plf
                plfs.append(plf)
            except libelf.ElfException as e:
                LOG.error(f"error: '{path}'")
                LOG.error(e)

    return static, plfs


if __name__ == "__main__":
    makerel()