# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "pyelftools>=0.32",
# ]
# ///
from dataclasses import dataclass
from pathlib import Path
import struct
import sys
from elftools.elf.elffile import ELFFile 
from elftools.elf.relocation import RelocationSection 
from elftools.elf.sections import Section, Symbol, SymbolTableSection
from elftools.elf.constants import E_FLAGS, E_FLAGS_MASKS
from elftools.elf.enums import ENUM_RELOC_TYPE_MIPS
from elftools.elf.descriptions import _DESCR_RELOC_TYPE_MIPS
from enum import IntEnum

# Elftools enums are just dicts...
class MipsReloc(IntEnum):
    R_MIPS_NONE = 0
    R_MIPS_32 = 2
    R_MIPS_26 = 4
    R_MIPS_HI16 = 5
    R_MIPS_LO16 = 6
    R_MIPS_GPREL16 = 7
    R_MIPS_LITERAL = 8
    R_MIPS15_S3 = 119
    R_MIPS_DVP_11_PCREL = 120
    R_MIPS_DVP_27_S4 = 121
    R_MIPS_DVP_11_S4 = 122
    R_MIPS_DVP_U15_S3 = 123

_U32 = struct.Struct("<I")
_S16 = struct.Struct("<h")
_U16 = struct.Struct("<H")

def read_generic(f, _s: struct.Struct, pos: int = -1) -> int:
    if pos != -1:
        old = f.tell()
        f.seek(pos)
    value = _s.unpack(f.read(_s.size))[0]
    if pos != -1:
        f.seek(old)
    return value

def read_u32(f, pos: int = -1) -> int:
    return read_generic(f, _U32, pos)

def read_u16(f, pos: int = -1) -> int:
    return read_generic(f, _U16, pos)

def read_s16(f, pos: int = -1) -> int:
    return read_generic(f, _S16, pos)

@dataclass
class RelInfo:
    rom_off: int
    addr: int
    type: MipsReloc
    sym_name: str
    sym_addr: int
    addend: int = 0
    comment_out: bool = False


EF_MIPS_MACH = 0x00FF0000
E_MIPS_MACH_R5900 = 0x00920000

ELF_PATH = Path(r"config/SLUS_210.07")

# VADDR for the start of the .lit4/LITERAL section
LIT4_SECTION_START = 0x003C7C80

def main():
    _monkey_patch()
    
    with ELFFile(ELF_PATH.open("rb")) as elf:
        if not check_elf(elf):
            sys.exit(-1)

        make_rels_file(elf)
        make_syms_file(elf)


def make_syms_file(elf: ELFFile) -> None:
    nono_chars = [
        "<",
        ">",
        ":",
        "\"",
        "/",
        "\\",
        "|",
        "?",
        "*"
    ]

    known_ldsyms = {
        "__bss_start",
        "__bss_end",
        "__bss_size",
        "_fbss",
        "__data_start",
        "__data_end",
        "_heap_size",
        "_stack_size",
        "_end",
        "_gp",
        # Should these be here? Technically they are used by mwInit
        # "__exception_table_start__",
        # "__exception_table_end__",
        # "__static_init",
        # "__static_init_end",
        # "_overlay_group_addresses",
    }

    # There should only be one symtab anyway
    symtab = get_symtab_section(elf)

    if symtab is None:
        return

    with open("symbols.txt", "w+") as rf:   
        all_syms = list(symtab.iter_symbols())

        seen_symbols = set()
        seen_addr = set()
        duped_symbols = set()
        duped_addr = set()

        for sym in all_syms:
            name = sym.name
            addr = sym.entry["st_value"]
            if name in seen_symbols:
                duped_symbols.add(name)
            elif addr in seen_addr:
                duped_addr.add(addr)
            else:
                seen_symbols.add(name)
                seen_addr.add(addr)
        
        for sym in all_syms:
            addr = sym.entry["st_value"]
            size = sym.entry["st_size"]
            # vis seems to always be STV_DEFAULT
            vis = sym.entry["st_other"]["visibility"]
            # local seems to always be 0
            local = sym.entry["st_other"]["local"]
            bind = sym.entry["st_info"]["bind"]
            type = sym.entry["st_info"]["type"]

            # assert (vis == "STV_DEFAULT" and local == 0)

            # Skip common autogenerated symbols
            if sym.name in ("__gnu_compiled_c", "gcc2_compiled."):
                continue
            
            # Anything below 0x100000 shouldn't exist
            if addr < 0x100000:
                continue
            
            # Skip section symbols
            if type == "STT_SECTION":
                continue

            comment = ""

            # Skip local ASM labels, VU microcode labels and known linker symbols
            if (sym.name.startswith("$") 
                or sym.name.startswith("_$")
                or sym.name.startswith(".vif")
                or sym.name.startswith(".dma")
                or sym.name.startswith(".gif")
                or sym.name.startswith(".vu")
                or sym.name in known_ldsyms
                ):
                remove = "// "
            else:
                remove = ""

            if size != 0:
                comment += f" size:0x{size:X}"
            
            match (bind):
                case "STB_GLOBAL":
                    # The default
                    comment += " visibility:global"
                    pass
                case "STB_LOCAL":
                    comment += " visibility:local"
                case "STB_WEAK":
                    comment += " visibility:weak"
                case "STB_LOPROC":
                    # Custom MWCC thing, it's called "MULTIDEF"
                    # appears exactly once in Code Veronica
                    # due to `inline TYPE function();` declaration
                    # maps roughly to `weak`, but it's distinct
                    comment += " visibility:weak"
                case _:
                    raise ValueError(f"Unknown symbol binding {bind} ({sym.name})")
            
            if sym.name in duped_symbols or addr in duped_addr:
                comment += " allow_duplicated:True"
            if len(sym.name) > 128:
                comment += f" filename:long_symbol_{addr:08X}"
            elif any(char in sym.name for char in nono_chars):
                comment += f" filename:bad_symbol_{addr:08X}"
            
            rf.write(f"{remove}{sym.name:<50} = 0x{addr:08X}; //{comment}\n")


def make_rels_file(elf: ELFFile) -> None:        
   
    all_relocs = parse_rels(elf)

    with open("relocs.txt", "w+") as rf:    
        for reloc in all_relocs:
            addend = reloc.addend
            rom_off = reloc.rom_off

            if addend == 0:
                addend_str = ""
            else:
                sign = "+" if addend > 0 else "-"
                addend_str = f"addend:{sign}0x{abs(addend):X}"

            if reloc.comment_out:
                cmm = "// "
            else:
                cmm = ""

            rf.write(f"{cmm}rom:0x{rom_off:06X} reloc:{reloc.type.name[2:]:<19} {addend_str:<15} symbol:{reloc.sym_name}\n")

def get_reloc_section(elf: ELFFile) -> tuple[RelocationSection | None, SymbolTableSection | None]:
    relmain: RelocationSection | None = elf.get_section_by_name(".relmain")
    if relmain is None:
        print("No .relmain section!")
        return None, None
    
    if relmain['sh_link'] == 0:
        print("Reloc section has no associated symbol table!")
        sys.exit(-1)

    symtable: SymbolTableSection = elf.get_section(relmain['sh_link'])

    return relmain, symtable

def get_symtab_section(elf: ELFFile) -> tuple[SymbolTableSection | None]:
    symtable: SymbolTableSection | None = elf.get_section_by_name(".symtab")
    if symtable is None:
        print("No .symtab section!")
    
    return symtable

def parse_rels(elf: ELFFile) -> list[RelInfo]:
    # Factor this out?
    main_section: Section | None = elf.get_section_by_name("main")
    if main_section is None:
        print("No main section!")
        sys.exit(-1)
    
    vaddr_diff = main_section.header["sh_addr"] - main_section.header["sh_offset"]

    relmain, symtab = get_reloc_section(elf)

    if relmain is None or symtab is None:
        return list()

    gp_value = _get_gp_value(elf)
    
    pending_hi: list[RelInfo] = []
    seen_hi: dict[str, int] = {}
    all_relocs: list[RelInfo] = []
    for r in relmain.iter_relocations():
        s: Symbol = symtab.get_symbol(r['r_info_sym'])
        reloc = RelInfo(
            r["r_offset"] - vaddr_diff,
            0,
            MipsReloc(r["r_info_type"]),
            s.name,
            s.entry["st_value"]
        )
        
        instr = read_u32(elf.stream, reloc.rom_off)
        simm = read_s16(elf.stream, reloc.rom_off)
        uimm = instr & 0xFFFF
        
        match reloc.type:
            case MipsReloc.R_MIPS_NONE:
                # NONE-thing to do
                pass
            case MipsReloc.R_MIPS_32:
                reloc.addr = instr
                reloc.addend = instr - reloc.sym_addr
            case MipsReloc.R_MIPS_26:
                reloc.addr = (instr & 0x03FFFFFF) * 4
                reloc.addend = reloc.addr - reloc.sym_addr
            case MipsReloc.R_MIPS_HI16:
                reloc.addr = uimm << 0x10
                pending_hi.append(reloc)
            case MipsReloc.R_MIPS_LO16:
                # This logic roughly matches GNU AS (_bfd_mips_elf_hi16_reloc)
                # and produces corrected addends for all the "orphaned" LO16s
                # in Code Veronica save for one ambiguous entry at ROM:14AA8
                # so I guess it's good enough.
                for hi in pending_hi:
                    hi.addr += simm
                    hi.addend = hi.addr - hi.sym_addr
                    if simm < 0:
                        seen_hi[hi.sym_name] = (hi.addr + 0x10000) & 0xFFFF0000
                    else:
                        seen_hi[hi.sym_name] = hi.addr & 0xFFFF0000
                
                reloc.addr = seen_hi[reloc.sym_name] + simm
                reloc.addend = reloc.addr - reloc.sym_addr
                pending_hi.clear()
            case MipsReloc.R_MIPS_GPREL16:
                reloc.addr = gp_value + simm
                reloc.addend = reloc.addr - reloc.sym_addr
            case MipsReloc.R_MIPS_LITERAL:
                # Literals are also not supported, as GNU AS
                # doesn't have a reloc operator for them...
                # We can lie and say it's a gprel reloc instead but meh
                reloc.addend = gp_value + simm - LIT4_SECTION_START
                # reloc.type = MipsReloc.R_MIPS_GPREL16
                # reloc.sym_name = f"D_{gp_value + simm:08X}"
                reloc.comment_out = True
            case MipsReloc.R_MIPS15_S3:
                # Needs upstream support...
                reloc.comment_out = True
            case (MipsReloc.R_MIPS_DVP_U15_S3 | MipsReloc.R_MIPS_DVP_27_S4):
                # Nothing to do, it's only for VU1 microcode
                reloc.comment_out = True
            case _:
                raise ValueError(f"Unimplemented reloc {reloc.type}")
            
        all_relocs.append(reloc)
    return all_relocs

def _get_gp_value(elf: ELFFile):
    reginfo = elf.get_section_by_name(".reginfo")

    # It better not use GP
    if reginfo is None:
        return 0

    addr = reginfo.header["sh_offset"]
    return read_u32(elf.stream, addr + 0x14)

def check_elf(elf: ELFFile) -> bool:
    flags = elf.header["e_flags"]
    type = elf.header["e_type"]

    if type != "ET_EXEC":
        print("Not a static file!")
        return False


    if ((flags & E_FLAGS_MASKS.EFM_MIPS_ABI) == E_FLAGS_MASKS.EFM_MIPS_ABI_EABI64 
        and (flags & E_FLAGS.EF_MIPS_ARCH) != E_FLAGS.EF_MIPS_ARCH_3 
        and(flags & EF_MIPS_MACH) != E_MIPS_MACH_R5900):
        print("Not an R5900 elf file!")
        return False
    
    if elf.header["e_shnum"] == 0:
        print("No section headers!")
        return False

    return True

def _monkey_patch():
    # Add DVP reloc types, though splat doesn't support them as of yet
    # TODO: ask for R_MIPS15_S3 support at some point, the others
    # are only for the VU microcode, which is out of scope
    ENUM_RELOC_TYPE_MIPS["R_MIPS15_S3"] = 119
    ENUM_RELOC_TYPE_MIPS["R_MIPS_DVP_11_PCREL"] = 120
    ENUM_RELOC_TYPE_MIPS["R_MIPS_DVP_27_S4"] = 121
    ENUM_RELOC_TYPE_MIPS["R_MIPS_DVP_11_S4"] = 122
    ENUM_RELOC_TYPE_MIPS["R_MIPS_DVP_U15_S3"] = 123
    _DESCR_RELOC_TYPE_MIPS[119] = "R_MIPS15_S3"
    _DESCR_RELOC_TYPE_MIPS[120] = "R_MIPS_DVP_11_PCREL"
    _DESCR_RELOC_TYPE_MIPS[121] = "R_MIPS_DVP_27_S4"
    _DESCR_RELOC_TYPE_MIPS[122] = "R_MIPS_DVP_11_S4"
    _DESCR_RELOC_TYPE_MIPS[123] = "R_MIPS_DVP_U15_S3"

if __name__ == "__main__":
    main()