#include "idt.hpp"
#include "portio.hpp"
#include "textoutput.hpp"

#include <array>
#include <cstdint>
#include <utility>

extern TextOutput& term;

static constexpr std::uint8_t TaskGate   = 0x5;
static constexpr std::uint8_t IntrGate16 = 0x6;
static constexpr std::uint8_t TrapGate16 = 0x7;
static constexpr std::uint8_t IntrGate32 = 0xE;
static constexpr std::uint8_t TrapGate32 = 0xF;

struct idt_entry_bits {
    std::uint32_t offset_low             : 16;
    std::uint32_t segment_selector       : 16;
    std::uint32_t rsvd                   :  8;
    std::uint32_t gate_type              :  4;
    std::uint32_t rsvd2                  :  1;
    std::uint32_t dpl                    :  2;
    std::uint32_t present                :  1;
    std::uint32_t offset_high            : 16;
} __attribute__((packed));

static std::array<Callback, 48> callbacks;

extern "C"
void interruptGeneralHandler(Registers regs)
{
    const auto& inum = regs.inum;

    if (inum >= 32) {
        if (inum >= 40)
            outb(0xA0, 0x20);

        outb(0x20, 0x20);
    }

    if (inum < callbacks.size()) {
       if (auto cb = callbacks[inum]; cb)
           cb(regs);
    }
}

template<std::size_t N>
struct StubEntry
{
    static constexpr bool HasError = N == 8 || (N >= 10 && N <= 14) || N == 17 || N == 30;

    __attribute__((naked))
    static void stub() {
        if constexpr (!HasError)
            asm volatile("push $0x0");

        asm volatile(R"(
            pusha
            push %0
            cld
            call interruptGeneralHandler
            pop %%eax
            popa
            add $0x4, %%esp
            iret
        )" :: "i"(N));
    }

    static constexpr std::uint32_t segment(std::uint16_t gdt_idx, bool useLdt, std::uint16_t rpl) {
        return gdt_idx | (useLdt ? 0x4 : 0x0) | (rpl & 0x3);
    }

    idt_entry_bits entry = {
        .offset_low = (uint32_t)stub & 0xFFFF,
        .segment_selector = segment(0x8, false, 0),
        .gate_type = IntrGate32,
        .dpl = 0,
        .present = 1,
        .offset_high = (uint32_t)stub >> 16
    };

    operator idt_entry_bits() const noexcept {
        return entry;
    }
};

static auto idt =
    []<std::size_t... ints>(std::index_sequence<ints...>) {
        return std::array<idt_entry_bits, 256> { StubEntry<ints>()... };
    }(std::make_index_sequence<48>{});

void idt_initialize()
{
    auto idtr = reinterpret_cast<std::uint64_t>(idt.data());
    idtr <<= 16;
    idtr |= idt.size() * sizeof(idt[0]);

    asm volatile("lidt %0" :: "m"(idtr));
}

void idt_register_callback(std::size_t num, Callback cb)
{
    if (num < callbacks.size())
        callbacks[num] = cb;
}