/*******************************************************************************
* Copyright 2024 Intel Corporation
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
*     http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*******************************************************************************/

#include <cstdint>
#include <cstring>
#include <fstream>
#include <stdexcept>
#include <utility>
#include <vector>

#include "generator/microkernel/elf.hpp"
#include "gemmstone/microkernel/fuser.hpp"

GEMMSTONE_NAMESPACE_START
namespace microkernel {

static void fixupJumpTargets(uint8_t *start, size_t len, ptrdiff_t adjust);

void fuse(std::vector<uint8_t> &binary,
        const std::vector<uint8_t> &microkernel, long id) {
    auto base = binary.data();
    auto bytes = binary.size();

    auto fheaderPtr = reinterpret_cast<FileHeader *>(base);

    bool ok = bytes >= sizeof(fheaderPtr) && fheaderPtr->magic == ELFMagic
            && fheaderPtr->elfClass == ELFClass64
            && fheaderPtr->endian == ELFLittleEndian
            && fheaderPtr->sectionHeaderSize == sizeof(SectionHeader)
            && (fheaderPtr->version == 0 || fheaderPtr->version == ELFVersion1)
            && (fheaderPtr->type == ZebinExec
                    || fheaderPtr->type == ELFRelocatable)
            && bytes >= sizeof(fheaderPtr)
                            + sizeof(SectionHeader) * fheaderPtr->sectionCount;

    if (!ok)
        throw std::runtime_error(
                "IGC did not generate a valid zebin program binary");

    bool foundZeInfo = false;
    const char *snames = nullptr;
    std::vector<std::pair<SectionHeader *, int>> textSections;

    auto *sheaders = reinterpret_cast<SectionHeader *>(
            base + fheaderPtr->sectionTableOff);

    snames = reinterpret_cast<char *>(
            base + sheaders[fheaderPtr->strTableIndex].offset);

    for (int s = 0; s < fheaderPtr->sectionCount; s++) {
        switch (sheaders[s].type) {
            case SectionHeader::Type::ZeInfo: foundZeInfo = true; break;
            case SectionHeader::Type::Program: {
                if (snames) {
                    std::string sname(snames + sheaders[s].name);
                    if (sname == ".text.Intel_Symbol_Table_Void_Program")
                        continue;
                    if (sname.substr(0, 6) != ".text.") continue;
                }
                textSections.emplace_back(sheaders + s, s);
                break;
            }
            default: break;
        }
    }

    if (!foundZeInfo || textSections.empty())
        throw std::runtime_error(
                "IGC did not generate a valid zebin program binary");

    for (auto &entry : textSections) {
        auto *text = entry.first;
        int textSectionID = entry.second;
        if (text->offset + text->size > bytes) continue;

        auto *insn = reinterpret_cast<const uint32_t *>(base + text->offset);
        auto *iend = reinterpret_cast<const uint32_t *>(
                base + text->offset + text->size);

        const uint8_t *spliceStart = nullptr;
        const uint8_t *spliceEnd = nullptr;

        for (; insn < iend; insn += 4) {
            if (insn[0] & (1u << 29))
                insn -= 2;
            else if (insn[3] == (sigilStart ^ id))
                spliceStart = reinterpret_cast<const uint8_t *>(insn);
            else if (insn[3] == (sigilEnd ^ id)) {
                spliceEnd = reinterpret_cast<const uint8_t *>(insn);
                break;
            }
        }

        if (!spliceStart || !spliceEnd) continue;

        int relSectionID = -1;
        std::string rname = ".rel";
        rname += (snames + text->name);
        for (int s = 0; s < fheaderPtr->sectionCount; s++) {
            if (sheaders[s].type != SectionHeader::Type::Relocation) continue;
            if (rname != (snames + sheaders[s].name)) continue;
            if (relSectionID >= 0)
                throw std::runtime_error(
                        "Multiple relocation sections for kernel");
            relSectionID = s;
        }

        auto removeBytes = spliceEnd - spliceStart + 16;

        size_t before = spliceStart - base;
        auto after = bytes - before - removeBytes;
        ptrdiff_t sizeAdjust = microkernel.size() - removeBytes;

        auto kbefore = before - text->offset;
        auto kafter = text->size - kbefore - removeBytes;

        std::vector<uint8_t> newBinary(bytes + sizeAdjust);
        auto newBase = newBinary.data();

        memmove(newBase, base, before);
        memmove(newBase + before, microkernel.data(), microkernel.size());
        memmove(newBase + before + microkernel.size(),
                spliceStart + removeBytes, after);

        fixupJumpTargets(newBase + text->offset, kbefore, +sizeAdjust);
        fixupJumpTargets(
                newBase + before + microkernel.size(), kafter, -sizeAdjust);

        fheaderPtr = reinterpret_cast<FileHeader *>(newBase);

        if (fheaderPtr->sectionTableOff > before)
            fheaderPtr->sectionTableOff += sizeAdjust;

        sheaders = reinterpret_cast<SectionHeader *>(
                newBase + fheaderPtr->sectionTableOff);
        sheaders[textSectionID].size += sizeAdjust;
        for (int s = 0; s < fheaderPtr->sectionCount; s++)
            if (sheaders[s].offset > before) sheaders[s].offset += sizeAdjust;

        if (relSectionID >= 0) {
            auto relSection = sheaders + relSectionID;
            auto rel = reinterpret_cast<Relocation *>(
                    newBase + relSection->offset);
            auto relEnd = reinterpret_cast<Relocation *>(
                    newBase + relSection->offset + relSection->size);
            for (; rel < relEnd; rel++) {
                if (rel->offset >= kbefore) rel->offset += sizeAdjust;
            }
        }

#ifdef SPLICE_DEBUG
        std::ofstream dump0("original." + std::to_string(id) + ".bin");
        dump0.write((const char *)binary.data(), binary.size());

        std::ofstream dump("patched." + std::to_string(id) + ".bin");
        dump.write((const char *)newBinary.data(), newBinary.size());
#endif

        std::swap(binary, newBinary);

        // Tail-recurse to handle any further instances of this microkernel
        fuse(binary, microkernel, id);
        return;
    }
}

void fuse(std::vector<uint8_t> &binary, const char *source) {
    std::vector<uint8_t> microkernel;
    const auto sigilLen = strlen(sigilBinary);

    auto toNybble = [](char c) {
        return ((c >= 'A') ? (c - 'A' + 10) : (c - '0')) & 0xF;
    };

    for (const char *s = std::strstr(source, sigilBinary); s;
            s = std::strstr(s, sigilBinary)) {
        s += sigilLen;
        char *after;
        long id = strtol(s, &after, 10);
        microkernel.clear();
        for (s = after + 1; *s != '\n'; s += 2) {
            if (!s[0] || !s[1]) break;
            microkernel.push_back(static_cast<uint8_t>(
                    (toNybble(s[0]) << 4) | toNybble(s[1])));
        }
        fuse(binary, microkernel, id);
    }
}

static void fixupJumpTargets(uint8_t *start, size_t len, ptrdiff_t adjust) {
    auto istart = reinterpret_cast<int32_t *>(start);
    auto iend = reinterpret_cast<int32_t *>(start + len);

    for (auto insn = istart; insn < iend; insn += 4) {
        if (insn[0] & (1u << 29)) {
            insn -= 2; /* skip compacted instructions */
            continue;
        }
        uint8_t op = insn[0] & 0xFF;
        if ((op & 0xF0) != 0x20) continue; /* skip non-jumps */
        if (op == 0x2B || op == 0x2D) continue; /* skip ret/calla */
        bool hasUIP = (op == 0x22 || op == 0x23 || op == 0x24 || op == 0x28
                || op == 0x2A || op == 0x2E);

        auto jumpFixup = [=](int32_t &ip) {
            auto target = ((insn - istart) << 2) + ip;
            if (target < 0 || target >= ptrdiff_t(len))
                ip += static_cast<int32_t>(adjust);
        };

        if (hasUIP) jumpFixup(insn[2]);
        jumpFixup(insn[3]);
    }
}

bool hasMicrokernels(const char *source) {
    return std::strstr(source, sigilBinary);
}

}
GEMMSTONE_NAMESPACE_END
