#include "vm.hpp"
#include <array>
#include <cstdint>
#include <cstdlib>
#include <iostream>
#include <print>
#include <stdexcept>
#include <string_view>
#include <thread>
#include <utility>
#include <vector>

using namespace vc5;
using namespace std::literals;

void VM::load(uint16_t offset, const uint8_t* data, size_t data_size)
{
    std::memcpy(&m_mem[offset], data, data_size);
}

namespace {

class Debugger {
public:
    explicit Debugger(VM& vm)
        : vm(&vm)
    {
    }

    void run_on_line(bool halted)
    {
        if (halted) {
            m_halted_last_call = true;
            return;
        }

        if (m_continue) {
            if (m_halted_last_call and m_break_on_halt) {
                m_continue = false;
            } else {
                auto ip = vm->reg(Reg::Rip);
                for (const auto& addr : m_break_on_address) {
                    if (addr == ip) {
                        m_continue = false;
                        break;
                    }
                }
            }

            if (m_continue)
                return;
        }

        m_halted_last_call = false;

        auto line = std::string();
        do {
            std::print("> ");
            std::cout.flush();
            std::getline(std::cin, line);

            if (line == "r") {
                for (uint16_t reg = 0; reg < 10; ++reg) {
                    constexpr auto reg_strs = std::array {
                        "r0"sv,
                        "r1"sv,
                        "r2"sv,
                        "r3"sv,
                        "r4"sv,
                        "r5"sv,
                        "rbp"sv,
                        "rsp"sv,
                        "rfl"sv,
                        "rip"sv,
                    };
                    std::println("    {: <3} | {:04x} {: 5}",
                        reg_strs[reg],
                        vm->reg(static_cast<Reg>(reg)),
                        vm->reg(static_cast<Reg>(reg)));
                }
            } else if (line == "c") {
                m_continue = true;
                break;
            } else if (line == "brhlt") {
                m_break_on_halt = true;
            } else if (line.starts_with("p")) {
                uint16_t v
                    = std::strtoul(line.c_str() + 2, nullptr, 16) & 0xfffe;
                std::println("[{: 4x}] {: 4x}", v, vm->word(v));
            } else if (line.starts_with("br")) {
                uint16_t v
                    = std::strtoul(line.c_str() + 3, nullptr, 16) & 0xfffe;
                m_break_on_address.push_back(v);
            }
        } while (!line.empty());
    }

private:
    VM* vm;
    std::vector<uint16_t> m_break_on_address {};
    bool m_halted_last_call = false;
    bool m_break_on_halt = false;
    bool m_continue = false;
};

}

int VM::run()
{
    constexpr uint16_t bootloader_size = 512;
    for (uint16_t i = 0; i * m_disk->block_size < bootloader_size; ++i) {
        m_disk->read(&m_mem[i * m_disk->block_size], i);
    }

    auto debugger = Debugger(*this);

    m_on = true;
    while (m_on) {
        if (not m_halted) {
            int result = run_instruction();
            if (result != 0)
                return result;
        }

        // debugger.run_on_line(m_halted);

        poll_events();
        if (m_halted) {
            std::this_thread::sleep_for(15ms);
        }
    }
    return 0;
}

uint16_t VM::reg(Reg reg) const
{
    return m_regs[std::to_underlying(reg)];
}

uint16_t VM::word(uint16_t addr) const
{
    if ((addr & 0b1) != 0) {
        throw std::invalid_argument(
            std::format("invalid address alignment, addr = {}", addr));
    }

    uint16_t value = 0;
    value |= static_cast<uint16_t>(m_mem[addr]) << 8;
    value |= m_mem[addr + 1];
    return value;
}

uint8_t VM::byte(uint16_t addr) const
{
    return m_mem[addr];
}

void VM::set_reg(Reg reg, uint16_t value)
{
    m_regs[std::to_underlying(reg)] = value;
}

void VM::set_word(uint16_t addr, uint16_t value)
{
    if ((addr & 0b1) != 0) {
        throw std::invalid_argument(
            std::format("invalid address alignment, addr = {}", addr));
    }

    m_mem[addr] = value >> 8;
    m_mem[addr + 1] = value & 0xff;
}

void VM::set_byte(uint16_t addr, uint8_t value)
{
    m_mem[addr] = value;
}

uint16_t VM::eat_word()
{
    auto ins = word(*m_rip);
    *m_rip += 2;
    return ins;
}

namespace {

struct InsReader {
    InsReader(VM& vm)
        : vm(&vm)
        , ins(vm.eat_word())
    {
    }

    VM* vm;
    uint16_t ins;

    auto op() -> Op
    {
        return static_cast<Op>(ins & 0b11'1111);
    }

    auto to_u16() const -> uint16_t
    {
        return ins;
    }

    auto dst_reg() const -> Reg
    {
        return static_cast<Reg>(ins >> 13 & 0b111);
    }

    auto op1_reg() const -> Reg
    {
        return static_cast<Reg>(ins >> 10 & 0b111);
    }

    auto op2_reg() const -> Reg
    {
        return static_cast<Reg>(ins >> 7 & 0b111);
    }

    auto dst() const -> uint16_t
    {
        return vm->reg(dst_reg());
    }

    auto op1() const -> uint16_t
    {
        return vm->reg(op1_reg());
    }

    auto op2() const -> uint16_t
    {

        return vm->reg(op2_reg());
    }

    auto op1_or_imm() -> uint16_t
    {
        return reg_val_or_imm(6, 10, 0b111);
    }

    auto op2_or_imm() -> uint16_t
    {
        return reg_val_or_imm(6, 7, 0b111);
    }

    auto reg_val_or_imm(
        uint16_t is_imm_bit, uint16_t reg_bit, uint16_t reg_mask) -> uint16_t
    {
        bool is_imm = (ins >> is_imm_bit & 1) != 0;
        if (is_imm) {
            return vm->eat_word();
        } else {
            return vm->reg(static_cast<Reg>(ins >> reg_bit & reg_mask));
        }
    }

    auto mem_addr(uint16_t base, uint16_t imm) const -> uint16_t
    {
        bool is_imm = to_u16() >> 6 & 1;
        if (is_imm) {
            return imm;
        }

        auto offset = static_cast<int16_t>(imm);
        if (offset >= 0) {
            return base + static_cast<uint16_t>(offset);
        } else {
            return base - static_cast<uint16_t>(-offset);
        }
    }
};

}

int VM::run_instruction()
{
    auto ip = *m_rip;
    if (ip >= 0x400)
        return -1;

    auto ins = InsReader(*this);
    auto op = ins.op();

    std::println(
        "    [{: >3x}]: 0x{:04x} {: <6} {:04b} {:04b} {:04b} {:04b}    {:04x}",
        ip,
        ins.to_u16(),
        op_str(op),
        ins.to_u16() >> 12 & 0xf,
        ins.to_u16() >> 8 & 0xf,
        ins.to_u16() >> 4 & 0xf,
        ins.to_u16() & 0xf,
        word(*m_rip));

    switch (op) {
        case Op::Nop:
            break;
        case Op::Hlt:
            if (not m_interrupts_enabled) {
                std::println("warning: halted without interrupts enabled");
            }
            m_halted = true;
            break;
        case Op::Jmp: {
            auto op1 = ins.op1_or_imm();

            if (op1 == 1) {
                for (uint16_t reg = 0; reg < 10; ++reg) {
                    constexpr auto reg_strs = std::array {
                        "r0"sv,
                        "r1"sv,
                        "r2"sv,
                        "r3"sv,
                        "r4"sv,
                        "r5"sv,
                        "rbp"sv,
                        "rsp"sv,
                        "rfl"sv,
                        "rip"sv,
                    };
                    std::println("    {: <3} | {:04x} {: >5}",
                        reg_strs[reg],
                        this->reg(static_cast<Reg>(reg)),
                        this->reg(static_cast<Reg>(reg)));
                }
                break;
            }

            *m_rip = op1;
            break;
        }
        case Op::Jnz: {
            auto op1 = ins.op1();
            auto op2 = ins.op2_or_imm();

            if (op1 != 0) {
                *m_rip = op2;
            }
            break;
        }
        case Op::Mov: {
            auto dst = static_cast<Reg>(ins.to_u16() >> 12 & 0b1111);
            auto src = ins.reg_val_or_imm(6, 7, 0b1111);
            set_reg(dst, src);
            break;
        }
        case Op::LdW: {
            auto addr = ins.mem_addr(ins.op1(), eat_word());

            if ((addr & 0b1) != 0) {
                std::println(stderr,
                    "error: invalid address alignment, addr = {}",
                    addr);
                *m_rfl |= 1 << std::to_underlying(Flag::Err);
                m_on = false;
                goto halt_execution;
            }

            auto reg = ins.dst_reg();
            set_reg(reg, word(addr));
            break;
        }
        case Op::StW: {
            auto addr = ins.mem_addr(ins.dst(), eat_word());

            if ((addr & 0b1) != 0) {
                std::println(stderr,
                    "error: invalid address alignment, addr = {}",
                    addr);
                *m_rfl |= 1 << std::to_underlying(Flag::Err);
                m_on = false;
                goto halt_execution;
            }
            auto src = ins.op2();

            if (m_vcd_loaded) {
                vcd_maybe_send_byte(addr, src >> 8);
                vcd_maybe_send_byte(addr + 1, src & 0xff);
            }

            set_word(addr, src);

            break;
        }
        case Op::LdB: {
            auto addr = ins.mem_addr(ins.op1(), eat_word());
            auto reg = ins.dst_reg();

            set_reg(reg, byte(addr));
            break;
        }
        case Op::StB: {
            auto addr = ins.mem_addr(ins.dst(), eat_word());
            auto src = ins.op2();

            if (m_vcd_loaded) {
                vcd_maybe_send_byte(addr, src & 0xff);
            }

            m_mem[addr] = src & 0xff;
            break;
        }
        case Op::Cmp: {
            auto op1 = ins.op1();
            auto op2 = ins.op2_or_imm();

            if (op1 == op2) {
                *m_rfl |= 1u << std::to_underlying(Flag::Eq);
            } else {
                *m_rfl &= (uint16_t)~(1u << std::to_underlying(Flag::Eq));
            }
            if (op1 < op2) {
                *m_rfl |= 1u << std::to_underlying(Flag::Be);
            } else {
                *m_rfl &= (uint16_t)~(1u << std::to_underlying(Flag::Be));
            }
            if (op1 > op2) {
                *m_rfl |= 1u << std::to_underlying(Flag::Ab);
            } else {
                *m_rfl &= (uint16_t)~(1u << std::to_underlying(Flag::Ab));
            }
            if ((int16_t)op1 < (int16_t)op2) {
                *m_rfl |= 1u << std::to_underlying(Flag::Lt);
            } else {
                *m_rfl &= (uint16_t)~(1u << std::to_underlying(Flag::Lt));
            }
            if ((int16_t)op1 > (int16_t)op2) {
                *m_rfl |= 1u << std::to_underlying(Flag::Gt);
            } else {
                *m_rfl &= (uint16_t)~(1u << std::to_underlying(Flag::Gt));
            }
            break;
        }
        case Op::Or:
        case Op::And:
        case Op::Xor:
        case Op::Shl:
        case Op::RShl:
        case Op::Shr:
        case Op::RShr:
        case Op::Add:
        case Op::Sub:
        case Op::RSub:
        case Op::Mul:
        case Op::IMul: {
            auto op1 = ins.op1();
            auto op2 = ins.op2_or_imm();
            auto dst_reg = ins.dst_reg();

            uint16_t* dst = &m_regs[std::to_underlying(dst_reg)];

            auto as_u16 = [](auto v) { return static_cast<uint16_t>(v); };
            auto as_i16 = [](auto v) { return static_cast<int16_t>(v); };

            switch (op) {
                case Op::Or:
                    *dst = op1 | op2;
                    break;
                case Op::Xor:
                    *dst = op1 ^ op2;
                    break;
                case Op::And:
                    *dst = op1 & op2;
                    break;
                case Op::Shl:
                    *dst = as_u16(op1 << op2);
                    break;
                case Op::RShl:
                    *dst = as_u16(op2 << op1);
                    break;
                case Op::Shr:
                    *dst = as_u16(op1 >> op2);
                    break;
                case Op::RShr:
                    *dst = as_u16(op2 >> op1);
                    break;
                case Op::Add:
                    *dst = op1 + op2;
                    break;
                case Op::Sub:
                    *dst = op1 - op2;
                    break;
                case Op::RSub:
                    *dst = op2 - op1;
                    break;
                case Op::Mul:
                    *dst = op1 * op2;
                    break;
                case Op::IMul:
                    *dst = as_u16(as_i16(op1) * as_i16(op2));
                    break;
                default:
                    break;
            }

            break;
        }
        case Op::Call: {
            auto op1 = ins.op1_or_imm();

            auto return_addr = *m_rip;
            set_word(*m_rsp, return_addr);
            *m_rsp += 2;

            *m_rip = op1;
            break;
        }
        case Op::Ret: {
            *m_rsp -= 2;
            auto return_addr = word(*m_rsp);
            *m_rip = return_addr;
            break;
        }
        case Op::RetI: {
            *m_rsp -= 2;
            auto return_addr = word(*m_rsp);
            *m_rip = return_addr;

            m_interrupts_enabled = true;
            break;
        }
        case Op::LVCD: {
            auto op1 = ins.op1_or_imm();
            m_vcd_descriptor = VcdDescript {
                .map_base_addr = word(op1),
            };
            m_vcd_loaded = true;

            break;
        }
        case Op::LKBD: {
            auto op1 = ins.op1_or_imm();
            m_kbd_descriptor = KbdDescript {
                .status_addr = word(op1),
                .keycode_addr = word(op1 + 2),
                .press_int_handler = word(op1 + 4),
                .release_int_handler = word(op1 + 6),
            };
            m_kbd_loaded = true;
            break;
        }
        case Op::DSKR: {
            auto addr = ins.dst();
            auto block = ins.op1();
            m_disk->read(&m_mem[addr], block);
            break;
        }
        case Op::DSKW: {
            auto addr = ins.op1();
            auto block = ins.dst();
            m_disk->write(&m_mem[addr], block);
            break;
        }
    }

halt_execution:
    return 0;
}

void VM::poll_events()
{
    while (true) {
        auto event = m_device->poll_event();

        if (not event)
            break;

        if (event->type() == IoEvent::Type::Quit) {
            m_on = false;
            break;
        }

        if (event->type() == IoEvent::Type::KeyPress
            or event->type() == IoEvent::Type::KeyRelease) {
            if (m_kbd_loaded) {
                auto& d = m_kbd_descriptor;

                uint16_t flags = 0;
                if (event->type() == IoEvent::Type::KeyRelease) {
                    flags |= 1 << std::to_underlying(KbdStatusFlag::Type);
                }
                if (event->as_key_event().shift_pressed) {
                    flags |= 1 << std::to_underlying(KbdStatusFlag::Shift);
                }
                if (event->as_key_event().ctrl_pressed) {
                    flags |= 1 << std::to_underlying(KbdStatusFlag::Ctrl);
                }
                if (event->as_key_event().alt_pressed) {
                    flags |= 1 << std::to_underlying(KbdStatusFlag::Alt);
                }
                if (event->as_key_event().altgr_pressed) {
                    flags |= 1 << std::to_underlying(KbdStatusFlag::AltGr);
                }

                set_word(d.status_addr, flags);
                set_word(d.keycode_addr, event->as_key_event().key);
                if (event->type() == IoEvent::Type::KeyPress) {
                    if (d.press_int_handler) {
                        interrupt(d.press_int_handler);
                    }
                } else {
                    if (d.release_int_handler) {
                        interrupt(d.release_int_handler);
                    }
                }
            }
            continue;
        }
    }
}

void VM::vcd_maybe_send_byte(uint16_t addr, uint8_t value)
{
    if (not m_vcd_loaded)
        return;

    constexpr uint16_t screen_size = 40 * 24;
    const auto& d = m_vcd_descriptor;

    if (addr >= d.map_base_addr and addr < d.map_base_addr + screen_size) {

        m_device->set_char(addr - d.map_base_addr, value).value();
    }
}

void VM::interrupt(uint16_t handler_addr)
{
    if (not m_interrupts_enabled)
        return;

    set_word(*m_rsp, *m_rip);
    *m_rsp += 2;

    *m_rip = handler_addr;

    m_halted = false;
    m_interrupts_enabled = false;
}

auto vc5::op_str(Op op) -> std::string_view
{
    switch (op) {
        case Op::Nop:
            return "Nop";
        case Op::Hlt:
            return "Hlt";
        case Op::Jmp:
            return "Jmp";
        case Op::Jnz:
            return "Jnz";
        case Op::Mov:
            return "Mov";
        case Op::LdW:
            return "LdW";
        case Op::StW:
            return "StW";
        case Op::LdB:
            return "LdB";
        case Op::StB:
            return "StB";
        case Op::Cmp:
            return "Cmp";
        case Op::Or:
            return "Or";
        case Op::And:
            return "And";
        case Op::Xor:
            return "Xor";
        case Op::Shl:
            return "Shl";
        case Op::RShl:
            return "RShl";
        case Op::Shr:
            return "Shr";
        case Op::RShr:
            return "RShr";
        case Op::Add:
            return "Add";
        case Op::Sub:
            return "Sub";
        case Op::RSub:
            return "RSub";
        case Op::Mul:
            return "Mul";
        case Op::IMul:
            return "IMul";
        case Op::Call:
            return "Call";
        case Op::Ret:
            return "Ret";
        case Op::RetI:
            return "RetI";
        case Op::LVCD:
            return "LVCD";
        case Op::LKBD:
            return "LKBD";
        case Op::DSKR:
            return "DSKR";
        case Op::DSKW:
            return "DSKW";
    }
    return ">.<";
}