#!/usr/bin/env python3

# SPDX-FileCopyrightText: © 2022-2024 Decompollaborate
# SPDX-License-Identifier: MIT

from __future__ import annotations

import argparse
from pathlib import Path
import sys

from .. import common
from .. import elf32
from .. import mips
from .. import frontendCommon as fec

from .. import __version__

PROGNAME = "elfObjDisasm"


def getToolDescription() -> str:
    return "Experimental MIPS elf disassembler"

def addOptionsToParser(parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
    parser.add_argument("-V", "--version", action="version", version=f"%(prog)s {__version__}")

    parser.add_argument("binary", help="Path to input elf binary file")
    parser.add_argument("output", help="Path to output. Use '-' to print to stdout instead")

    parser.add_argument("--data-output", help="Path to output the data and rodata disassembly")

    parser.add_argument("--file-splits", help="Path to a file splits csv")

    parser.add_argument("--split-functions", help="Enables the function and rodata splitter. Expects a path to place the splited functions", metavar="PATH")

    parser.add_argument("--instr-category", help="The instruction category to use when disassembling every passed instruction. Defaults to 'cpu'", choices=["cpu", "rsp", "r3000gte", "r4000allegrex", "r5900"])

    parser.add_argument("--function-info", help="Specifies a path where to output a csvs sumary file of every analyzed function", metavar="PATH")


    readelfOptions = parser.add_argument_group("readelf-like flags")

    readelfOptions.add_argument("-a", "--all", help="Equivalent to --file-header --section-headers --syms --dyn-syms --relocs --display-got", action="store_true")
    readelfOptions.add_argument("--file-header", help="Display the ELF file header", action="store_true")
    readelfOptions.add_argument("-S", "--section-headers", help="Display the sections' header", action="store_true")
    readelfOptions.add_argument("-s", "--syms", help="Display the symbol table", action="store_true")
    readelfOptions.add_argument("--dyn-syms", help="Display the dynamic symbol table", action="store_true")
    readelfOptions.add_argument("-r", "--relocs", help="Display the relocations (if present)", action="store_true")
    readelfOptions.add_argument("--display-got", help="Shows Global offset table information", action="store_true")

    readelfOptions.add_argument("--readelf-only", help="Exit after processing the readelf-like flags, without performing any disassembly", action="store_true")


    common.Context.addParametersToArgParse(parser)

    common.GlobalConfig.addParametersToArgParse(parser)

    mips.InstructionConfig.addParametersToArgParse(parser)

    return parser

def getArgsParser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=getToolDescription(), prog=PROGNAME, formatter_class=common.Utils.PreserveWhiteSpaceWrapRawTextHelpFormatter)
    return addOptionsToParser(parser)

def applyArgs(args: argparse.Namespace) -> None:
    if args.libultra_syms is None:
        args.libultra_syms = False
    if args.hardware_regs is None:
        args.hardware_regs = False
    common.GlobalConfig.parseArgs(args)
    mips.InstructionConfig.parseArgs(args)

def applyGlobalConfigurations() -> None:
    common.GlobalConfig.REMOVE_POINTERS = False
    common.GlobalConfig.IGNORE_BRANCHES = False
    common.GlobalConfig.SYMBOL_FINDER_FILTER_LOW_ADDRESSES = False

    common.GlobalConfig.ALLOW_UNKSEGMENT = False

    common.GlobalConfig.INPUT_FILE_TYPE = common.InputFileType.ELF


def applyReadelfLikeFlags(elfFile: elf32.Elf32File, args: argparse.Namespace) -> None:
    if args.all:
        elfFile.readelf_fileHeader()
        elfFile.readelf_sectionHeaders()
        elfFile.readelf_syms()
        elfFile.readelf_dyn_syms()
        elfFile.readelf_relocs()
        elfFile.readelf_displayGot()
    else:
        if args.file_header:
            elfFile.readelf_fileHeader()

        if args.section_headers:
            elfFile.readelf_sectionHeaders()

        if args.syms:
            elfFile.readelf_syms()

        if args.dyn_syms:
            elfFile.readelf_dyn_syms()

        if args.relocs:
            elfFile.readelf_relocs()

        if args.display_got:
            elfFile.readelf_displayGot()

    if args.readelf_only:
        sys.exit(0)

SpecialSectionNames = {".text", ".data", ".rodata", ".bss"}

def getOutputPath(inputPath: Path, textOutput: Path, dataOutput: Path, sectionType: common.FileSectionType, sectionName: str) -> Path:
    outputPath = dataOutput
    if sectionType == common.FileSectionType.Text:
        outputPath = textOutput

    outputFilePath = outputPath
    if str(outputPath) != "-":
        outname = f"{inputPath.stem}"
        if sectionName not in SpecialSectionNames:
            outname += f"_{sectionName}"
        outputFilePath = outputFilePath / outname

    return outputFilePath

def processSection(
        context: common.Context, array_of_bytes: bytes,
        processedSections: dict[common.FileSectionType, list[mips.sections.SectionBase]],
        segmentPaths: dict[common.FileSectionType, list[Path]],
        sectionsPerName: dict[str, mips.sections.SectionBase],
        inputPath: Path, textOutput: Path, dataOutput: Path,
        sectionName: str, sectionEntry: elf32.Elf32SectionHeaderEntry,
        sectionType: common.FileSectionType,
        sectionClass: type[mips.sections.SectionText | mips.sections.SectionData | mips.sections.SectionRodata | mips.sections.SectionBss]
    ) -> None:
    outputFilePath = getOutputPath(inputPath, textOutput, dataOutput, sectionType, sectionName)

    vromStart = sectionEntry.offset
    vromEnd = vromStart + sectionEntry.size
    vramStart = sectionEntry.addr
    if sectionEntry.addr == 0:
        vramStart = sectionEntry.offset

    mipsSection: mips.sections.SectionBase


    inname = f"{inputPath.stem}"
    if sectionName not in SpecialSectionNames:
        inname += f"_{sectionName}"

    if issubclass(sectionClass, mips.sections.SectionBss):
        bssStart = vramStart
        bssEnd = bssStart + sectionEntry.size
        mipsSection = sectionClass(context, vromStart, vromEnd, bssStart, bssEnd, inname, 0, None)
    else:
        mipsSection = sectionClass(context, vromStart, vromEnd, vramStart, inname, array_of_bytes, 0, None)

    mipsSection.setCommentOffset(vromStart)
    mipsSection.customSectionName = sectionName
    processedSections[sectionType].append(mipsSection)
    segmentPaths[sectionType].append(outputFilePath)
    sectionsPerName[sectionName] = mipsSection


def getProcessedSections(context: common.Context, elfFile: elf32.Elf32File, array_of_bytes: bytes, inputPath: Path, textOutput: Path, dataOutput: Path) -> tuple[dict[common.FileSectionType, list[mips.sections.SectionBase]], dict[common.FileSectionType, list[Path]], dict[str, mips.sections.SectionBase]]:
    processedSections: dict[common.FileSectionType, list[mips.sections.SectionBase]] = {
        common.FileSectionType.Text: [],
        common.FileSectionType.Data: [],
        common.FileSectionType.Rodata: [],
        common.FileSectionType.Bss: [],
    }
    segmentPaths: dict[common.FileSectionType, list[Path]] = {
        common.FileSectionType.Text: [],
        common.FileSectionType.Data: [],
        common.FileSectionType.Rodata: [],
        common.FileSectionType.Bss: [],
    }
    sectionsPerName: dict[str, mips.sections.SectionBase] = dict()

    for sectionName, sectionEntry in elfFile.progbitsExecute.items():
        processSection(context, array_of_bytes, processedSections, segmentPaths, sectionsPerName, inputPath, textOutput, dataOutput, sectionName, sectionEntry, common.FileSectionType.Text, mips.sections.SectionText)

    for sectionName, sectionEntry in elfFile.progbitsWrite.items():
        processSection(context, array_of_bytes, processedSections, segmentPaths, sectionsPerName, inputPath, textOutput, dataOutput, sectionName, sectionEntry, common.FileSectionType.Data, mips.sections.SectionData)

    for sectionName, sectionEntry in elfFile.progbitsNoWrite.items():
        processSection(context, array_of_bytes, processedSections, segmentPaths, sectionsPerName, inputPath, textOutput, dataOutput, sectionName, sectionEntry, common.FileSectionType.Rodata, mips.sections.SectionRodata)

    for sectionName, sectionEntry in elfFile.nobitsPerName.items():
        processSection(context, array_of_bytes, processedSections, segmentPaths, sectionsPerName, inputPath, textOutput, dataOutput, sectionName, sectionEntry, common.FileSectionType.Bss, mips.sections.SectionBss)

    return processedSections, segmentPaths, sectionsPerName

def changeGlobalSegmentRanges(context: common.Context, processedSegments: dict[common.FileSectionType, list[mips.sections.SectionBase]]) -> None:
    lowestVromStart: int|None = None
    highestVromEnd: int|None = None
    lowestVramStart: int|None = None
    highestVramEnd: int|None = None

    for subSegment in processedSegments.values():
        for section in subSegment:
            if lowestVromStart is None or section.vromStart < lowestVromStart:
                lowestVromStart = section.vromStart
            if highestVromEnd is None or section.vromEnd > highestVromEnd:
                highestVromEnd = section.vromEnd
            if lowestVramStart is None or section.vram < lowestVramStart:
                lowestVramStart = section.vram
            if highestVramEnd is None or section.vramEnd > highestVramEnd:
                highestVramEnd = section.vramEnd

    if lowestVromStart is None:
        lowestVromStart = 0x0
    if highestVromEnd is None:
        highestVromEnd = 0xFFFFFFFF
    if lowestVramStart is None:
        lowestVramStart = 0x0
    if highestVramEnd is None:
        highestVramEnd = 0xFFFFFFFF
    context.changeGlobalSegmentRanges(lowestVromStart, highestVromEnd, lowestVramStart, highestVramEnd)


def addContextSymFromSymEntry(context: common.Context, symEntry: elf32.Elf32SymEntry, symAddress: int, symName: str|None=None, symVrom: int|None=None) -> common.ContextSymbol|None:
    if symAddress == 0:
        return None

    segment = context.globalSegment
    if not segment.isVramInRange(symAddress):
        segment = context.unknownSegment

    if symEntry.stType == elf32.Elf32SymbolTableType.FUNC.value:
        segment = context.globalSegment
        contextSym = segment.addFunction(symAddress, vromAddress=symVrom)
    elif symEntry.stType == elf32.Elf32SymbolTableType.OBJECT.value:
        contextSym = segment.addSymbol(symAddress, vromAddress=symVrom)
    elif symEntry.stType == elf32.Elf32SymbolTableType.SECTION.value:
        # print(symEntry)
        return None
    elif symEntry.stType == elf32.Elf32SymbolTableType.NOTYPE.value:
        if symEntry.shndx == elf32.Elf32SectionHeaderNumber.ABS.value:
            segment = context.globalSegment
            contextSym = segment.addSymbol(symAddress, vromAddress=symVrom)
            contextSym.isElfNotype = True
        else:
            bind = elf32.Elf32SymbolTableBinding.fromValue(symEntry.stBind)
            if bind != elf32.Elf32SymbolTableBinding.LOCAL:
                common.Utils.eprint(f"Warning: Non-LOCAL ({bind}) NOTYPE symbol '{symName}' has an unhandled shndx value: '0x{symEntry.shndx:X}'")
            contextSym = segment.addSymbol(symAddress, vromAddress=symVrom)
    else:
        common.Utils.eprint(f"Warning: symbol '{symName}' has an unhandled stType: '{symEntry.stType}'")
        contextSym = segment.addSymbol(symAddress, vromAddress=symVrom)

    if symName is not None:
        if symName.startswith("."):
            contextSym._isStatic = True
        else:
            contextSym.name = symName
    contextSym.isUserDeclared = True
    contextSym.setSizeIfUnset(symEntry.size)

    bind = elf32.Elf32SymbolTableBinding.fromValue(symEntry.stBind)
    if bind is not None:
        contextSym.visibility = bind.name.lower()

    return contextSym


def insertSymtabIntoContext(context: common.Context, symbolTable: elf32.Elf32Syms, stringTable: elf32.Elf32StringTable, elfFile: elf32.Elf32File) -> None:
    # Use the symbol table to replace symbol names present in disassembled sections
    for i, symEntry in enumerate(symbolTable):
        symName = stringTable[symEntry.name]

        if symEntry.shndx == 0:
            continue

        sectHeaderEntry = elfFile.sectionHeaders[symEntry.shndx]
        if sectHeaderEntry is None:
            continue

        if elfFile.header.type != elf32.Elf32ObjectFileType.REL.value:
            if symEntry.value != 0:
                addContextSymFromSymEntry(context, symEntry, symEntry.value, symName)
            continue

        sectName = elfFile.shstrtab[sectHeaderEntry.name]
        sectType = common.FileSectionType.fromStr(sectName)
        if sectType == common.FileSectionType.Invalid:
            common.Utils.eprint(f"Warning: symbol {i} (name: '{symName}', value: 0x{symEntry.value:X}) is referencing invalid section '{sectName}'")
            continue

        symVrom = symEntry.value + sectHeaderEntry.offset
        symAddress = symVrom

        contextSym = addContextSymFromSymEntry(context, symEntry, symAddress, symName, symVrom)
        if contextSym is not None:
            contextSym.sectionType = sectType

def insertDynsymIntoContext(context: common.Context, symbolTable: elf32.Elf32Syms, stringTable: elf32.Elf32StringTable) -> None:
    for symEntry in symbolTable:
        if symEntry.value == 0 or symEntry.shndx == 0:
            continue

        symName = stringTable[symEntry.name]
        addContextSymFromSymEntry(context, symEntry, symEntry.value, symName)

def insertGotIntoContext(context: common.Context, got: elf32.Elf32GlobalOffsetTable, stringTable: elf32.Elf32StringTable) -> None:
    lazyResolver = got.localsTable[0]
    lazyResolverSym = context.globalSegment.addSymbol(lazyResolver)
    lazyResolverSym.name = "$$.LazyResolver"
    lazyResolverSym.isUserDeclared = True
    lazyResolverSym.isGotLocal = True

    gotIndex = len(got.localsTable)

    for gotEntry in got.globalsTable:
        symName = stringTable[gotEntry.symEntry.name]

        contextSym = addContextSymFromSymEntry(context, gotEntry.symEntry, gotEntry.getAddress(), symName)
        if contextSym is not None:
            contextSym.gotIndex = gotIndex
        gotIndex += 1


def injectAllElfSymbols(context: common.Context, elfFile: elf32.Elf32File, processedSegments: dict[common.FileSectionType, list[mips.sections.SectionBase]], sectionsPerName: dict[str, mips.sections.SectionBase]) -> None:
    if elfFile.symtab is not None and elfFile.strtab is not None:
        # Inject symbols from the reloc table referenced in each section
        if elfFile.header.type == elf32.Elf32ObjectFileType.REL.value:
            for sectionName, relocs in elfFile.relPerName.items():
                subSegment = sectionsPerName.get(sectionName, None)
                for rel in relocs:
                    symbolEntry = elfFile.symtab[rel.rSym]
                    symbolName = elfFile.strtab[symbolEntry.name]

                    if symbolEntry.stType != elf32.Elf32SymbolTableType.SECTION.value:
                        if symbolName == "":
                            continue

                    if subSegment is None:
                        common.Utils.eprint(f"Warning: reloc '{rel}' references unhandled section '{sectionName}'")
                        continue

                    relocVrom = subSegment.vromStart + rel.offset
                    relocInfo = context.addGlobalReloc(relocVrom, common.RelocType(rel.rType), symbolName)
                    if symbolEntry.stType == elf32.Elf32SymbolTableType.SECTION.value:
                        sectionEntry = elfFile.sectionHeaders[symbolEntry.shndx]
                        assert sectionEntry is not None, rel

                        sectionType = common.FileSectionType.fromStr(elfFile.shstrtab[sectionEntry.name])
                        if sectionType != common.FileSectionType.Invalid:
                            sectionVram = processedSegments[sectionType][0].vram
                            relocInfo.staticReference = common.RelocationStaticReference(sectionType, sectionVram)

        # Use the symtab to replace symbol names present in disassembled sections
        insertSymtabIntoContext(context, elfFile.symtab, elfFile.strtab, elfFile)

    if elfFile.dynsym is not None and elfFile.dynstr is not None:
        # Use the dynsym to replace symbol names present in disassembled sections
        insertDynsymIntoContext(context, elfFile.dynsym, elfFile.dynstr)

    if elfFile.got is not None and elfFile.dynstr is not None:
        insertGotIntoContext(context, elfFile.got, elfFile.dynstr)

def processGlobalOffsetTable(context: common.Context, elfFile: elf32.Elf32File) -> None:
    if elfFile.dynamic is not None:
        common.GlobalConfig.GP_VALUE = elfFile.dynamic.getGpValue()

    if elfFile.reginfo is not None:
        common.GlobalConfig.GP_VALUE = elfFile.reginfo.gpValue

    if elfFile.dynamic is not None and elfFile.got is not None:
        localsTable = [common.GotEntry(address, False) for address in elfFile.got.localsTable]
        globalsTable = [common.GotEntry(gotEntry.getAddress(), True) for gotEntry in elfFile.got.globalsTable]
        tableEntries = localsTable + globalsTable

        assert elfFile.dynamic.pltGot is not None
        context.setupGotTable(elfFile.dynamic.pltGot, tableEntries)

        for small in elfFile.smallSections.values():
            context.addSmallSection(small.addr, small.size)


def processArguments(args: argparse.Namespace) -> int:
    applyArgs(args)

    applyGlobalConfigurations()

    context = common.Context()
    context.parseArgs(args)

    common.Utils.printQuietless(f"{PROGNAME} (spimdisasm {__version__})")

    inputPath = Path(args.binary)
    array_of_bytes = common.Utils.readFileAsBytearray(inputPath)
    elfFile = elf32.Elf32File(array_of_bytes)

    elfFile.handleHeaderIdent()
    elfFile.handleFlags()

    instrCategory = args.instr_category
    if instrCategory is None:
        if elf32.Elf32Constants.Elf32HeaderFlag._5900 in elfFile.elfFlags:
            instrCategory = "r5900"

    common.Utils.printQuietless(f"{PROGNAME} {inputPath}: Processing global offset table...")
    processGlobalOffsetTable(context, elfFile)

    applyReadelfLikeFlags(elfFile, args)

    textOutput = Path(args.output)
    if args.data_output is None:
        dataOutput = textOutput
    else:
        dataOutput = Path(args.data_output)

    common.Utils.printQuietless(f"{PROGNAME} {inputPath}: Reading segments...")
    if args.file_splits is not None:
        splits = common.FileSplitFormat()
        splits.readCsvFile(Path(args.file_splits))
        processedSegments, segmentPaths = fec.FrontendUtilities.getSplittedSections(context, splits, array_of_bytes, inputPath, textOutput, dataOutput)
        sectionsPerName: dict[str, mips.sections.SectionBase] = dict()
        for sectType, subSection in processedSegments.items():
            if len(subSection) < 1:
                continue
            if sectType == common.FileSectionType.Text:
                sectionsPerName[".text"] = subSection[0]
            elif sectType == common.FileSectionType.Data:
                sectionsPerName[".data"] = subSection[0]
            elif sectType == common.FileSectionType.Rodata:
                sectionsPerName[".rodata"] = subSection[0]
            elif sectType == common.FileSectionType.Bss:
                sectionsPerName[".bss"] = subSection[0]
    else:
        processedSegments, segmentPaths, sectionsPerName = getProcessedSections(context, elfFile, array_of_bytes, inputPath, textOutput, dataOutput)

    changeGlobalSegmentRanges(context, processedSegments)

    fec.FrontendUtilities.configureProcessedFiles(processedSegments, instrCategory)

    common.Utils.printQuietless(f"{PROGNAME} {inputPath}: Injecting elf symbols...")
    injectAllElfSymbols(context, elfFile, processedSegments, sectionsPerName)

    processedFilesCount = 0
    for sect in processedSegments.values():
        processedFilesCount += len(sect)

    common.Utils.printQuietless(f"{PROGNAME} {inputPath}: Analyzing sections...")
    fec.FrontendUtilities.analyzeProcessedFiles(processedSegments, segmentPaths, processedFilesCount)

    common.Utils.printQuietless(f"{PROGNAME} {inputPath}: Writing files...")
    fec.FrontendUtilities.writeProcessedFiles(processedSegments, segmentPaths, processedFilesCount)

    if args.split_functions is not None:
        common.Utils.printQuietless(f"{PROGNAME} {inputPath}: Migrating functions and rodata...")
        functionMigrationPath = Path(args.split_functions)
        fec.FrontendUtilities.migrateFunctions(processedSegments, functionMigrationPath)

        common.Utils.printQuietless(f"{PROGNAME} {inputPath}: Generating functions list...")
        mips.FilesHandlers.writeMigratedFunctionsList(processedSegments, functionMigrationPath, inputPath.stem)

    if args.save_context is not None:
        common.Utils.printQuietless(f"{PROGNAME} {inputPath}: Writing context...")
        contextPath = Path(args.save_context)
        contextPath.parent.mkdir(parents=True, exist_ok=True)
        context.saveContextToFile(contextPath)

    if args.function_info is not None:
        fec.FrontendUtilities.writeFunctionInfoCsv(processedSegments, Path(args.function_info))

    common.Utils.printQuietless(f"{PROGNAME} {inputPath}: Done!")

    return 0

def addSubparser(subparser: argparse._SubParsersAction[argparse.ArgumentParser]) -> None:
    parser = subparser.add_parser(PROGNAME, help=getToolDescription(), formatter_class=common.Utils.PreserveWhiteSpaceWrapRawTextHelpFormatter)

    addOptionsToParser(parser)

    parser.set_defaults(func=processArguments)


def elfObjDisasmMain() -> int:
    args = getArgsParser().parse_args()

    return processArguments(args)
