#include <codegen/x86/X64InstructionRewriter.hpp>
#include <codegen/x86/X64InstructionDecoder.hpp>
#include <codegen/x86/X64OpcodeConstants.hpp>
#include <codegen/CodegenException.hpp>
#include <codegen/IInstructionScanner.hpp>
#include <cstdint>
#include <limits>

namespace Codegen {

using namespace X64OpcodeConstants;

X64InstructionRewriter::ReferenceSite X64InstructionRewriter::ClassifyInstruction(
    const std::uint8_t* data,
    std::size_t length
) {
    const auto info = X64InstructionDecoder{}.DecodeInstruction(data, length);

    if (info.HasRipRelativeDisp) {
        return {ReferenceKind::RipRelativeDisp32, info.RipRelativeDispOffset};
    }

    if (!info.HasBranchTarget) {
        return {ReferenceKind::None, 0};
    }

    const bool rel32 = info.IsTwoByteOpcode || info.Opcode == OneByteCallRel32 || info.Opcode == OneByteJmpRel32;
    return rel32
        ? ReferenceSite{ReferenceKind::Rel32Branch, info.Length - Disp32Size}
        : ReferenceSite{ReferenceKind::Rel8Branch, info.Length - Disp8Size};
}

std::int32_t X64InstructionRewriter::ReadInt32(const std::uint8_t* data) {
    return static_cast<std::int32_t>(
        static_cast<std::uint32_t>(data[0]) |
        (static_cast<std::uint32_t>(data[1]) << 8) |
        (static_cast<std::uint32_t>(data[2]) << 16) |
        (static_cast<std::uint32_t>(data[3]) << 24));
}

void X64InstructionRewriter::WriteInt32(std::uint8_t* data, std::int32_t value) {
    const auto unsignedValue = static_cast<std::uint32_t>(value);
    data[0] = static_cast<std::uint8_t>(unsignedValue & 0xFF);
    data[1] = static_cast<std::uint8_t>((unsignedValue >> 8) & 0xFF);
    data[2] = static_cast<std::uint8_t>((unsignedValue >> 16) & 0xFF);
    data[3] = static_cast<std::uint8_t>((unsignedValue >> 24) & 0xFF);
}

std::int64_t X64InstructionRewriter::AdjustedDelta(
    std::int64_t instrPosition,
    std::int64_t targetPosition,
    std::int64_t cutStart,
    std::int64_t cutEnd,
    std::int64_t sliceSize,
    std::int64_t delta
) {
    const bool instrShifts = instrPosition >= cutEnd && instrPosition < sliceSize;
    const bool targetShifts = targetPosition >= cutEnd && targetPosition < sliceSize;

    if (instrShifts == targetShifts) {
        return 0;
    }

    return instrShifts ? -delta : delta;
}

RewriteResult X64InstructionRewriter::Rewrite(
    const std::vector<std::uint8_t>& codeSection,
    const RewriteRequest& request
) const {
    const auto offset = static_cast<std::size_t>(request.Offset);

    if (offset > codeSection.size()) {
        throw CodegenException("Rewrite offset is out of bounds", request.Offset);
    }

    const X64InstructionDecoder decoder;
    const std::size_t decodedLength = decoder.Decode(codeSection.data() + offset, codeSection.size() - offset);

    const auto oldLength = static_cast<std::int64_t>(decodedLength);
    const auto newLength = static_cast<std::int64_t>(request.NewBytes.size());
    const std::int64_t delta = newLength - oldLength;

    const auto cutStart = static_cast<std::int64_t>(offset);
    const std::int64_t cutEnd = cutStart + oldLength;

    std::vector<std::uint8_t> rewritten;
    rewritten.reserve(codeSection.size() + static_cast<std::size_t>(delta > 0 ? delta : 0));
    rewritten.insert(rewritten.end(), codeSection.begin(), codeSection.begin() + offset);
    rewritten.insert(rewritten.end(), request.NewBytes.begin(), request.NewBytes.end());
    rewritten.insert(rewritten.end(), codeSection.begin() + offset + decodedLength, codeSection.end());

    const auto scanner = MakeInstructionScanner();
    const auto matches = scanner->ScanCodeSection(codeSection, 0, codeSection.size());

    std::vector<AddressAdjustment> adjustments;

    for (const auto& match : matches) {
        const auto matchOffset = static_cast<std::int64_t>(match.Offset);

        if (matchOffset >= cutStart && matchOffset < cutEnd) {
            continue;
        }

        const auto site = ClassifyInstruction(codeSection.data() + match.Offset, match.Length);

        if (site.Kind != ReferenceKind::None) {
            const std::int64_t instrEnd = matchOffset + static_cast<std::int64_t>(match.Length);
            const std::size_t fieldPosition = static_cast<std::size_t>(match.Offset) + site.FieldOffset;

            if (site.Kind == ReferenceKind::Rel8Branch) {
                const auto disp = static_cast<std::int8_t>(codeSection[fieldPosition]);
                const std::int64_t targetPosition = instrEnd + disp;
                const std::int64_t adjust = AdjustedDelta(
                    matchOffset, targetPosition, cutStart, cutEnd,
                    static_cast<std::int64_t>(codeSection.size()), delta);

                if (adjust != 0) {
                    const std::int64_t newDisp = static_cast<std::int64_t>(disp) + adjust;

                    if (newDisp < std::numeric_limits<std::int8_t>::min() ||
                        newDisp > std::numeric_limits<std::int8_t>::max()) {
                        throw CodegenException("Rel8 branch displacement overflow after rewrite", match.Offset);
                    }

                    const std::int64_t newMatchOffset = matchOffset < cutStart ? matchOffset : matchOffset + delta;
                    const auto newFieldPosition = static_cast<std::size_t>(
                        newMatchOffset + static_cast<std::int64_t>(site.FieldOffset));
                    rewritten[newFieldPosition] = static_cast<std::uint8_t>(static_cast<std::int8_t>(newDisp));
                }
            } else {
                const auto disp = ReadInt32(codeSection.data() + fieldPosition);
                const std::int64_t targetPosition = instrEnd + disp;
                const std::int64_t adjust = AdjustedDelta(
                    matchOffset, targetPosition, cutStart, cutEnd,
                    static_cast<std::int64_t>(codeSection.size()), delta);

                if (adjust != 0) {
                    const std::int64_t newDisp = static_cast<std::int64_t>(disp) + adjust;

                    if (newDisp < std::numeric_limits<std::int32_t>::min() ||
                        newDisp > std::numeric_limits<std::int32_t>::max()) {
                        throw CodegenException("Rel32/disp32 displacement overflow after rewrite", match.Offset);
                    }

                    const std::int64_t newMatchOffset = matchOffset < cutStart ? matchOffset : matchOffset + delta;
                    const auto newFieldPosition = static_cast<std::size_t>(
                        newMatchOffset + static_cast<std::int64_t>(site.FieldOffset));
                    WriteInt32(rewritten.data() + newFieldPosition, static_cast<std::int32_t>(newDisp));
                }
            }
        }

        if (matchOffset >= cutEnd && delta != 0) {
            const auto newOffset = static_cast<Domain::FileByteOffset>(matchOffset + delta);
            adjustments.push_back({match.Offset, newOffset, delta});
        }
    }

    return {rewritten, adjustments};
}

}