#include <domain/Types.hpp>
#include <io/FileReader.hpp>
#include <io/FileWriter.hpp>
#include <elfpatcher/linux/LinuxElfPatcher.hpp>
#include <elfpatcher/general/SegmentFilter.hpp>
#include <elfpatcher/general/EntryStubBuilder.hpp>
#include <elfpatcher/general/ProgramHeaderLayoutBuilder.hpp>
#include <elfpatcher/general/SectionHeaderTableBuilder.hpp>
#include <io/ByteWriter.hpp>
#include <relinker/parsing/ElfReader.hpp>
#include <relinker/analysis/ValidationPolicy.hpp>
#include <relinker/analysis/SyscallScanner.hpp>
#include <relinker/analysis/CallSiteResolver.hpp>
#include <relinker/analysis/UnusedNidFilter.hpp>
#include <relinker/output/SysVDynamicSectionBuilder.hpp>
#include <relinker/output/CallRegistryWriter.hpp>
#include <relinker/pipeline/RelinkerPipeline.hpp>
#include <codegen/IAmd64OnlyConverter.hpp>
#include <filesystem>
#include <iostream>
#include <memory>
#include <string>
#include <optional>
#include <elfpatcher/windows/WindowsElfPatcher.hpp>

int main(const int argc, char* argv[]) {
    bool skipSyscallCheck = false;
    bool toIntel = false;
    bool skipUnusedNidFilter = false;
    bool writeRegistry = false;
    bool toWindows = false;

    std::string inputPath;
    std::string outputPath;
    std::string runPath = "$ORIGIN/libs";

    for (int i = 1; i < argc; ++i) {
        const std::string arg = argv[i];
        if (arg == "--skip-syscall-check") {
            skipSyscallCheck = true;
        } else if (arg == "--to-intel") {
            toIntel = true;
        } else if (arg == "--skip-unused-nid-filter") {
            skipUnusedNidFilter = true;
        } else if (arg == "--registry") {
            writeRegistry = true;
        } else if (arg == "--rpath") {
            if (i + 1 >= argc) {
                std::cerr << "FAIL: --rpath requires a value\n";
                return 1;
            }
            runPath = argv[++i];
        } else if (arg == "--windows") {
            toWindows = true;
        } else if (inputPath.empty()) {
            inputPath = arg;
        } else if (outputPath.empty()) {
            outputPath = arg;
        } else {
            std::cerr << "FAIL: unexpected argument: " << arg << "\n";
            return 1;
        }
    }

    if (inputPath.empty() || outputPath.empty()) {
        std::cerr << "Usage: relinker [--windows] [--skip-syscall-check] [--to-intel] [--skip-unused-nid-filter] [--registry] [--rpath <path>] <input.elf> <output.elf>\n"
             "Example: relinker input.elf output.elf\n";
        return 1;
    }

    try {
        Io::FileReader fileReader;
        Io::FileWriter fileWriter;

        auto sourceBytes = fileReader.Read(inputPath);

        if (toIntel) {
            const Relinker::ElfReader elfReader(sourceBytes);
            const auto converter = Codegen::MakeAmd64OnlyConverter();
            auto result = converter->Convert(std::move(sourceBytes), elfReader.ReadCodeSegments());
            fileWriter.Write(outputPath, std::move(result.Bytes));
            std::cout << "OK: " << result.ReplacedCount << " instructions replaced\n";
            return 0;
        }

        auto elfReader = std::make_shared<Relinker::ElfReader>(sourceBytes);

        auto syscallScanner = skipSyscallCheck
            ? Relinker::MakeNullSyscallScanner()
            : Relinker::MakeSyscallScanner();

        const auto pipeline = std::make_shared<Relinker::RelinkerPipeline>(
            elfReader,
            std::move(syscallScanner),
            Relinker::MakeCallSiteResolver(),
            std::make_shared<Relinker::ValidationPolicy>(),
            std::make_shared<Relinker::SysVDynamicSectionBuilder>(),
            Relinker::MakeUnusedNidFilter(),
            !skipUnusedNidFilter && !toWindows
        );

        auto result = pipeline->Relink(sourceBytes);

        if (writeRegistry) {
            const std::filesystem::path outFsPath(outputPath);
            const std::string registryPath = (outFsPath.parent_path() / (outFsPath.stem().string() + ".registry.json")).string();
            auto callRegistryWriter = std::make_shared<Relinker::CallRegistryWriter>();
            fileWriter.Write(registryPath, callRegistryWriter->WriteCallRegistry(result.RegistryEntries));
        }

        auto byteWriter = std::make_shared<Io::ByteWriter>();

        std::shared_ptr<Elfpatcher::IElfPatcher> patcher;
        if (toWindows) {
            patcher = std::make_shared<Elfpatcher::Windows::WindowsPePatcher>();
        } else {
            patcher = std::make_shared<Elfpatcher::Linux::LinuxElfPatcher>(
                std::make_shared<Elfpatcher::EntryStubBuilder>(),
                std::make_shared<Elfpatcher::ProgramHeaderLayoutBuilder>(
                    std::make_shared<Elfpatcher::SegmentFilter>(),
                    byteWriter
                ),
                std::make_shared<Elfpatcher::SectionHeaderTableBuilder>(byteWriter),
                byteWriter
            );
        }
        auto patched = patcher->Patch(sourceBytes, result.OriginalHeaders, result.DynamicSection, result.OriginalPltGotVaddr, runPath);
        fileWriter.Write(outputPath, patched);

    } catch (const Domain::RelinkerException& e) {
        std::cerr << "FAIL: " << e.what();
        if (e.FailureOffset != 0)
            std::cerr << " (offset 0x" << std::hex << e.FailureOffset << ")";
        std::cerr << "\n";
        return 2;
    } catch (const std::exception& e) {
        std::cerr << "FAIL: " << e.what() << "\n";
        return 2;
    }

    return 0;
}