#include <elfpatcher/linux/LinuxElfPatcher.hpp>
#include <elfpatcher/general/EntryStubBuilder.hpp>
#include <elfpatcher/general/ProgramHeaderLayoutBuilder.hpp>
#include <elfpatcher/general/SectionHeaderTableBuilder.hpp>
#include <elfpatcher/general/SegmentFilter.hpp>
#include <io/ByteWriter.hpp>
#include <algorithm>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <map>
#include <memory>
#include <stdexcept>
#include <string>
#include <vector>

namespace {

using Bytes = std::vector<std::uint8_t>;

constexpr std::uint64_t kPage = 0x1000;

void Require(const bool ok, const std::string& what) {
    if (!ok) throw std::runtime_error("linux load pages: " + what);
}

template <typename T>
T Read(const Bytes& bytes, const std::uint64_t offset) {
    T value{};
    std::memcpy(&value, bytes.data() + offset, sizeof(T));
    return value;
}

template <typename T>
void Write(Bytes& bytes, const std::uint64_t offset, const T value) {
    std::memcpy(bytes.data() + offset, &value, sizeof(T));
}

Bytes Fixture(const std::size_t size) {
    Bytes bytes(size);
    bytes[0] = 0x7F; bytes[1] = 'E'; bytes[2] = 'L'; bytes[3] = 'F';
    bytes[4] = 2; bytes[5] = 1; bytes[6] = 1;
    Write<std::uint16_t>(bytes, 16, 3);
    Write<std::uint16_t>(bytes, 18, 62);
    Write<std::uint64_t>(bytes, 24, 0x400);
    Write<std::uint64_t>(bytes, 32, 64);
    Write<std::uint16_t>(bytes, 54, 56);
    Write<std::uint16_t>(bytes, 56, 1);
    bytes[0x400] = 0xC3;
    return bytes;
}

Domain::ProgramHeader Load(const std::uint64_t offset, const std::uint64_t vaddr, const std::uint64_t fileSize, const std::uint64_t memorySize, const std::uint32_t flags) {
    return {1, flags, offset, vaddr, vaddr, fileSize, memorySize, 0x1000};
}

Elfpatcher::Linux::LinuxElfPatcher Patcher() {
    const auto byteWriter = std::make_shared<Io::ByteWriter>();
    return 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);
}

std::vector<Domain::ProgramHeader> OutputLoads(const Bytes& output) {
    std::vector<Domain::ProgramHeader> loads;
    const auto table = Read<std::uint64_t>(output, 32);
    const auto count = Read<std::uint16_t>(output, 56);
    for (std::uint16_t index = 0; index < count; ++index) {
        const auto header = table + index * 56u;
        if (Read<std::uint32_t>(output, header) != 1) continue;
        loads.push_back({1, Read<std::uint32_t>(output, header + 4), Read<std::uint64_t>(output, header + 8), Read<std::uint64_t>(output, header + 16),
            Read<std::uint64_t>(output, header + 24), Read<std::uint64_t>(output, header + 32), Read<std::uint64_t>(output, header + 40), Read<std::uint64_t>(output, header + 48)});
    }
    return loads;
}

std::map<std::uint64_t, std::uint8_t> MapLikeLinux(const Bytes& output) {
    std::map<std::uint64_t, std::uint8_t> memory;
    for (const auto& load : OutputLoads(output)) {
        const auto prefix = load.MappedAddress % kPage;
        Require(load.Offset % kPage == prefix, "PT_LOAD offset and address differ within a page");
        const auto fileEnd = load.MappedAddress + load.FileSize;
        const auto filePageEnd = (fileEnd + kPage - 1) & ~(kPage - 1);
        for (auto address = load.MappedAddress - prefix; address < filePageEnd; ++address) {
            const auto offset = load.Offset - prefix + (address - (load.MappedAddress - prefix));
            memory[address] = offset < output.size() ? output[offset] : 0;
        }
        if (load.MemorySize > load.FileSize) {
            const auto memoryPageEnd = (load.MappedAddress + load.MemorySize + kPage - 1) & ~(kPage - 1);
            for (auto address = fileEnd; address < memoryPageEnd; ++address)
                memory[address] = 0;
        }
    }
    return memory;
}

void RequireImage(const Bytes& output, const Bytes& source, const std::vector<Domain::ProgramHeader>& headers, const std::string& what) {
    const auto memory = MapLikeLinux(output);
    for (const auto& load : headers) {
        for (std::uint64_t relative = 0; relative < load.MemorySize; ++relative) {
            const auto address = load.MappedAddress + relative;
            const auto found = memory.find(address);
            Require(found != memory.end(), what + ": address not mapped");
            const std::uint8_t expected = relative < load.FileSize ? source[load.Offset + relative] : 0;
            if (found->second != expected) {
                char text[96];
                std::snprintf(text, sizeof(text), ": byte at 0x%llx reads 0x%02x, not 0x%02x", static_cast<unsigned long long>(address), found->second, expected);
                Require(false, what + text);
            }
        }
    }
}

std::uint64_t OutputOffset(const Bytes& output, const std::uint64_t vaddr) {
    for (const auto& load : OutputLoads(output))
        if (load.MappedAddress == vaddr) return load.Offset;
    throw std::runtime_error("linux load pages: no output PT_LOAD at the segment's address");
}

void ZeroFillBelowTheNextSegment() {
    auto source = Fixture(0x3000);
    std::fill(source.begin() + 0x1000, source.begin() + 0x1100, std::uint8_t{0xAA});
    std::fill(source.begin() + 0x1100, source.begin() + 0x1c20, std::uint8_t{0x55});
    std::fill(source.begin() + 0x1c20, source.begin() + 0x2000, std::uint8_t{0x77});
    const std::vector<Domain::ProgramHeader> headers{
        Load(0x400, 0x400, 0x10, 0x10, 5),
        Load(0x1000, 0x1000, 0x100, 0x1c20, 6),
        Load(0x1c20, 0x2c20, 0x3e0, 0x3e0, 0),
    };
    const auto output = Patcher().Patch(source, headers, {}, 0, "$ORIGIN/libs", true, false, {});
    RequireImage(output, source, headers, "bss sharing a page with the next segment");
}

void DataAndZeroFillBelowTheNextSegment() {
    auto source = Fixture(0x4000);
    std::fill(source.begin() + 0x1000, source.begin() + 0x2400, std::uint8_t{0xAA});
    std::fill(source.begin() + 0x2400, source.begin() + 0x2c20, std::uint8_t{0x55});
    std::fill(source.begin() + 0x2c20, source.begin() + 0x3000, std::uint8_t{0x77});
    const std::vector<Domain::ProgramHeader> headers{
        Load(0x400, 0x400, 0x10, 0x10, 5),
        Load(0x1000, 0x1000, 0x1400, 0x1c20, 6),
        Load(0x2c20, 0x2c20, 0x3e0, 0x3e0, 0),
    };
    const auto output = Patcher().Patch(source, headers, {}, 0, "$ORIGIN/libs", true, false, {});
    RequireImage(output, source, headers, "data and bss sharing a page with the next segment");
}

void ZeroFileBytesStayInPlace() {
    auto source = Fixture(0x3000);
    std::fill(source.begin() + 0x1000, source.begin() + 0x1100, std::uint8_t{0xAA});
    std::fill(source.begin() + 0x2c20, source.begin() + 0x3000, std::uint8_t{0x77});
    const std::vector<Domain::ProgramHeader> headers{
        Load(0x400, 0x400, 0x10, 0x10, 5),
        Load(0x1000, 0x1000, 0x100, 0x1c20, 6),
        Load(0x2c20, 0x2c20, 0x3e0, 0x3e0, 0),
    };
    const auto output = Patcher().Patch(source, headers, {}, 0, "$ORIGIN/libs", true, false, {});
    RequireImage(output, source, headers, "zero bytes below the next segment");
    Require(OutputOffset(output, 0x2c20) == 0x2c20, "a segment whose page already reads right was moved");
}

}

int main() {
    try {
        ZeroFillBelowTheNextSegment();
        DataAndZeroFillBelowTheNextSegment();
        ZeroFileBytesStayInPlace();
    } catch (const std::exception& error) {
        std::cerr << error.what() << '\n';
        return 1;
    }
    return 0;
}