#include "emulator/hart.h"

#include <stdio.h>
#include <inttypes.h>
#include <assert.h>
#include <stdbool.h>
#include <string.h>

#define REG_ZERO 0 
#define REG_RA   1 
#define REG_SP   2 
#define REG_GP   3 
#define REG_TP   4 
#define REG_T0   5 
#define REG_T1   6 
#define REG_T2   7 
#define REG_S0   8 
#define REG_S1   9 
#define REG_A0   10 
#define REG_A1   11 
#define REG_A2   12 
#define REG_A3   13 
#define REG_A4   14 
#define REG_A5   15 
#define REG_A6   16 
#define REG_A7   17 
#define REG_S2   18 
#define REG_S3   19 
#define REG_S4   20 
#define REG_S5   21 
#define REG_S6   22 
#define REG_S7   23 
#define REG_S8   24 
#define REG_S9   25 
#define REG_S10  26 
#define REG_S11  27 
#define REG_T3   28 
#define REG_T4   29 
#define REG_T5   30 
#define REG_T6   31 

static inline uint32_t sign_extend(uint32_t word, uint32_t size)
{
    const uint32_t mask = 1U << (size - 1);
    return (word ^ mask) - mask;
}

struct Instruction
{
    uint8_t  opcode;
    uint8_t  rs1;
    uint8_t  rs2;
    uint8_t  rd;
    uint8_t  funct3;
    uint8_t  funct7;
    uint32_t imm;
};
typedef struct Instruction Instruction;

static Instruction decode_r_type(uint32_t word)
{
    Instruction instruction = {0};
    instruction.opcode =  word        & 0x7F;
    instruction.rd     = (word >> 7)  & 0x1F;
    instruction.funct3 = (word >> 12) & 0x07;
    instruction.rs1    = (word >> 15) & 0x1F;
    instruction.rs2    = (word >> 20) & 0x1F;
    instruction.funct7 =  word >> 25;
    return instruction;
};

static Instruction decode_i_type(uint32_t word)
{
    Instruction instruction = {0};
    instruction.opcode =  word        & 0x7F;
    instruction.rd     = (word >> 7)  & 0x1F;
    instruction.funct3 = (word >> 12) & 0x07;
    instruction.rs1    = (word >> 15) & 0x1F;
    instruction.imm    = sign_extend(word >> 20, 12);
    return instruction;
};

static Instruction decode_s_type(uint32_t word)
{
    Instruction instruction = {0};
    instruction.opcode =  word        & 0x7F;
    instruction.funct3 = (word >> 12) & 0x07;
    instruction.rs1    = (word >> 15) & 0x1F;
    instruction.rs2    = (word >> 20) & 0x1F;
    instruction.imm    = sign_extend(((word >> 7) & 0x1F) | (word >> 25), 12);
    return instruction;
};

static Instruction decode_b_type(uint32_t word)
{
    Instruction instruction = decode_s_type(word);
    instruction.imm = ((instruction.imm << 11) & 0x800) | (instruction.imm & 0xfffff7ff);
    return instruction;
};

static Instruction decode_u_type(uint32_t word)
{
    Instruction instruction = {0};
    instruction.opcode =  word        & 0x7F;
    instruction.rd     = (word >> 7)  & 0x1F;
    instruction.imm    =  word        & 0xFFFFF000;
    return instruction;
};

static Instruction decode_j_type(uint32_t word)
{
    Instruction instruction = {0};
    instruction.opcode =  word        & 0x7F;
    instruction.rd     = (word >> 7)  & 0x1F;
    instruction.imm    = sign_extend(
        ((word & 0x80000000) >> 11) |
        ((word & 0x000FF000) >> 0)  |
        ((word & 0x00100000) >> 9)  |
        ((word & 0x7FE00000) >> 20), 21);
    return instruction;
}

static void execute_op_imm(Hart* hart, uint32_t instruction)
{
    const Instruction inst = decode_i_type(instruction);
    if (inst.rd == 0) return;
    
    switch (inst.funct3)
    {
        case 0: // ADDI
            hart->regs[inst.rd] = hart->regs[inst.rs1] + inst.imm;
            break;
        case 1: // SLLI
            hart->regs[inst.rd] = hart->regs[inst.rs1] << (inst.imm & 0x1F);
            break;
        case 2: // SLTI
            hart->regs[inst.rd] = (int32_t)hart->regs[inst.rs1] < (int32_t)inst.imm ? 1 : 0;
            break;
        case 3: // SLTIU
            hart->regs[inst.rd] = hart->regs[inst.rs1] < inst.imm ? 1 : 0;
            break;
        case 4: // XORI
            hart->regs[inst.rd] = hart->regs[inst.rs1] ^ inst.imm;
            break;
        case 5: // SRLI, SRAI
        {
            const uint32_t shamt = inst.imm & 0x1F;
            uint32_t res = hart->regs[inst.rs1] >> shamt;
            if ((inst.imm & 0x400) && shamt > 0) { res = sign_extend(res, 32 - shamt); }
            hart->regs[inst.rd] = res;
            break;
        }
        case 6: // ORI
            hart->regs[inst.rd] = hart->regs[inst.rs1] | inst.imm;
            break;
        case 7: // ANDI
            hart->regs[inst.rd] = hart->regs[inst.rs1] & inst.imm;
            break;
        default:
            assert(!"Unhandled OP-IMM");
    }
}

static void execute_op(Hart* hart, uint32_t instruction)
{
    const Instruction inst = decode_r_type(instruction);
    if (inst.rd == 0) return;
    
    switch (inst.funct3)
    {
        case 0: // ADD, SUB
            if (instruction & 0x40000000)
            {
                hart->regs[inst.rd] = hart->regs[inst.rs1] - hart->regs[inst.rs2];
            }
            else
            {
                hart->regs[inst.rd] = hart->regs[inst.rs1] + hart->regs[inst.rs2];
            }
            break;
        case 1: // SLL
            hart->regs[inst.rd] = hart->regs[inst.rs1] << (hart->regs[inst.rs2] & 0x1F);
            break;
        case 2: // SLT
            hart->regs[inst.rd] = (int32_t)hart->regs[inst.rs1] < (int32_t)hart->regs[inst.rs2] ? 1 : 0;
            break;
        case 3: // SLTU
            hart->regs[inst.rd] = hart->regs[inst.rs1] < hart->regs[inst.rs2] ? 1 : 0;
            break;
        case 4: // XOR
            hart->regs[inst.rd] = hart->regs[inst.rs1] ^ hart->regs[inst.rs2];
            break;
        case 5: // SRL, SRA
        {
            const uint32_t shamt = hart->regs[inst.rs2] & 0x1F;
            uint32_t res = hart->regs[inst.rs1] >> shamt;
            if ((instruction & 0x40000000) && shamt > 0) { res = sign_extend(res, 32 - shamt); }
            hart->regs[inst.rd] = res;
            break;
        }
        case 6: // OR
            hart->regs[inst.rd] = hart->regs[inst.rs1] | hart->regs[inst.rs2];
            break;
        case 7: // AND
            hart->regs[inst.rd] = hart->regs[inst.rs1] & hart->regs[inst.rs2];
            break;
        default:
            assert(!"Unhandled OP-IMM");
    }
}

static void execute_branch(Hart* hart, uint32_t instruction)
{
    const Instruction inst = decode_b_type(instruction);
    const uint32_t r1 = hart->regs[inst.rs1];
    const uint32_t r2 = hart->regs[inst.rs2];
    bool take_branch = false;

    switch (inst.funct3)
    {
        case 0: take_branch = (r1 == r2); break;
        case 1: take_branch = (r1 != r2); break;
        case 4: take_branch = (r1 < r2); break;
        case 5: take_branch = (r1 >= r2); break;
        case 6: take_branch = ((int32_t)r1 < (int32_t)r2); break;
        case 7: take_branch = ((int32_t)r1 >= (int32_t)r2); break;
    }

    if (take_branch)
    {
        hart->pc += inst.imm;
    }
    else
    {
        hart->pc += 4;
    }
}

static inline uint32_t load_size(Hart* hart, uint32_t address, uint32_t size)
{
    if ((address & 0x80000000) == 0)
    {
        assert(address + size < hart->mem_size);
        uint32_t value = 0;
        memcpy(&value, hart->mem + address, size);
        return value;
    }
    
    return 0;
}

static uint32_t load_byte(Hart* hart, uint32_t address)
{
    return load_size(hart, address, 1);
}

static uint32_t load_half(Hart* hart, uint32_t address)
{
    return load_size(hart, address, 2);
}

static uint32_t load_word(Hart* hart, uint32_t address)
{
    return load_size(hart, address, 4);
}

static inline void store_size(Hart* hart, uint32_t address, uint32_t value, uint32_t size)
{
    if ((address & 0x80000000) == 0)
    {
        assert(address + size < hart->mem_size);
        memcpy(hart->mem + address, &value, size);
    }
    else if (address == 0x80000000)
    {
        fwrite(&value, 1, size, stdout);
    }
}

static void store_byte(Hart* hart, uint32_t address, uint8_t value)
{
    store_size(hart, address, value, 1);
}

static void store_half(Hart* hart, uint32_t address, uint16_t value)
{
    store_size(hart, address, value, 2);
}

static void store_word(Hart* hart, uint32_t address, uint32_t value)
{
    store_size(hart, address, value, 4);
}

static void execute_op_load(Hart* hart, uint32_t instruction)
{
    const Instruction inst = decode_i_type(instruction);
    const uint32_t address = hart->regs[inst.rs1] + inst.imm;
    
    uint32_t value = 0;
    
    switch (inst.funct3 & 0x03)
    {
        case 0:
            value = load_byte(hart, address);
            if ((inst.funct3 & 0x40) == 0)
            {
                value = sign_extend(value, 8);
            }
            break;
        case 1:
            value = load_half(hart, address);
            if ((inst.funct3 & 0x40) == 0)
            {
                value = sign_extend(value, 16);
            }
            break;
        case 2:
            value = load_byte(hart, address);
            break;
        default:
            assert(!"Unhandled load size");
            break;
    }

    if (inst.rd != 0)
    {
        hart->regs[inst.rd] = value;
    }
}

static void execute_op_store(Hart* hart, uint32_t instruction)
{
    const Instruction inst = decode_s_type(instruction);
    const uint32_t address = hart->regs[inst.rs1] + inst.imm;
    const uint32_t value = hart->regs[inst.rs2];

    switch (inst.funct3 & 0x03)
    {
        case 0: store_byte(hart, address, value & 0xFF); break;
        case 1: store_half(hart, address, value & 0xFFFF); break;
        case 2: store_word(hart, address, value); break;
        default:
            assert(!"Unhandled store size");
            break;
    }
}

static void execute_misc_mem(Hart* hart, uint32_t instruction)
{
    const Instruction inst = decode_i_type(instruction);
    
    if (inst.funct3 == 0)
    {
        // FENCE
    }
    else
    {
        assert(!"Unhandled MISC-MEM instruction");
    }
}

static void handle_ecall(Hart* hart)
{
    if (hart->regs[REG_A0] == 0)
    {
        hart->halted = true;
    }
}

static uint32_t crs_read(Hart* hart, uint16_t id)
{
    // @Todo
    return 0;
}

static void crs_write(Hart* hart, uint16_t id, uint32_t new_value)
{
    // @Todo
}

static uint32_t crs_read_write(Hart* hart, uint16_t id, uint32_t new_value)
{
    // @Todo
    return 0;
}

static uint32_t crs_read_clear(Hart* hart, uint16_t id, uint32_t mask)
{
    // @Todo
    return 0;
}

static uint32_t crs_read_set(Hart* hart, uint16_t id, uint32_t mask)
{
    // @Todo
    return 0;
}

static void execute_system(Hart* hart, uint32_t instruction)
{
    const Instruction inst = decode_i_type(instruction);

    switch (inst.funct3)
    {
        case 0:
        {
            if (inst.rs1 == 0 && inst.rd == 0)
            {
                if (inst.imm == 0)
                {
                    handle_ecall(hart);
                }
                else if (inst.imm == 1)
                {
                    // EBREAK
                }
                else
                {
                    assert(!"Unhandled SYSTEM/PRIV instruction");
                }
            }
            else
            {
                assert(!"Unhandled SYSTEM/PRIV instruction");
            }
            break;
        }
        case 1: // CRSRW
        {
            uint16_t crs_id = inst.imm & 0xFFF;
            
            if (inst.rd == 0)
            {
                crs_write(hart, crs_id, hart->regs[inst.rs1]);
            }
            else
            {
                hart->regs[inst.rd] = crs_read_write(hart, crs_id, hart->regs[inst.rs1]);
            }
            break;
        }
        case 2: // CRSRS
        {
            uint16_t crs_id = inst.imm & 0xFFF;
            uint32_t value = 0;
            
            if (inst.rs1 == 0)
            {
                value = crs_read(hart, crs_id);
            }
            else
            {
                value = crs_read_set(hart, crs_id, hart->regs[inst.rs1]);
            }
            
            if (inst.rd != 0)
            {
                hart->regs[inst.rd] = value;
            }
            break;
        }
        case 3: // CRSRC
        {
            uint16_t crs_id = inst.imm & 0xFFF;
            uint32_t value = 0;
            
            if (inst.rs1 == 0)
            {
                value = crs_read(hart, crs_id);
            }
            else
            {
                value = crs_read_clear(hart, crs_id, hart->regs[inst.rs1]);
            }
            
            if (inst.rd != 0)
            {
                hart->regs[inst.rd] = value;
            }
            break;
        }
        case 5: // CRSRWI
        {
            uint16_t crs_id = inst.imm & 0xFFF;
            
            if (inst.rd == 0)
            {
                crs_write(hart, crs_id, inst.rs1);
            }
            else
            {
                hart->regs[inst.rd] = crs_read_write(hart, crs_id, inst.rs1);
            }
            break;
        }
        case 6: // CRSRSI
        {
            uint16_t crs_id = inst.imm & 0xFFF;
            uint32_t value = 0;
            
            if (inst.rs1 == 0)
            {
                value = crs_read(hart, crs_id);
            }
            else
            {
                value = crs_read_set(hart, crs_id, inst.rs1);
            }
            
            if (inst.rd != 0)
            {
                hart->regs[inst.rd] = value;
            }
            break;
        }
        case 7: // CRSRCI
        {
            uint16_t crs_id = inst.imm & 0xFFF;
            uint32_t value = 0;
            
            if (inst.rs1 == 0)
            {
                value = crs_read(hart, crs_id);
            }
            else
            {
                value = crs_read_clear(hart, crs_id, inst.rs1);
            }
            
            if (inst.rd != 0)
            {
                hart->regs[inst.rd] = value;
            }
            break;
        }
        default:
            assert(!"Unhandled SYSTEM instruction");
            break;
    }
}

void execute(Hart* hart, uint32_t instruction)
{
    switch (instruction & 0x7f)
    {
        case 0x03:
            execute_op_load(hart, instruction);
            hart->pc += 4;
            break;
        case 0x0F:
            execute_misc_mem(hart, instruction);
            hart->pc += 4;
            break;
        case 0x13:
            execute_op_imm(hart, instruction);
            hart->pc += 4;
            break;
        case 0x17: // AUIPC
        {
            Instruction inst = decode_u_type(instruction);
            if (inst.rd != 0)
            {
                hart->regs[inst.rd] = inst.imm + hart->pc;
            }
            hart->pc += 4;
            break;
        }
        case 0x23:
            execute_op_store(hart, instruction);
            hart->pc += 4;
            break;
        case 0x33:
            execute_op(hart, instruction);
            hart->pc += 4;
            break;
        case 0x37: // LUI
        {
            Instruction inst = decode_u_type(instruction);
            if (inst.rd != 0)
            {
                hart->regs[inst.rd] = inst.imm;
            }
            hart->pc += 4;
            break;
        }
        case 0x63:
            execute_branch(hart, instruction);
            break;
        case 0x67: // JALR
        {
            Instruction inst = decode_i_type(instruction);
            assert(inst.funct3 == 0);
            if (inst.rd != 0)
            {
                hart->regs[inst.rd] = hart->pc + 4;
            }
            hart->pc = (hart->regs[inst.rs1] + inst.imm) & 0xFFFFFFFE;
            break;
        }
        case 0x6F: // JAL
        {
            Instruction inst = decode_j_type(instruction);
            if (inst.rd != 0)
            {
                hart->regs[inst.rd] = hart->pc + 4;
            }
            hart->pc += inst.imm;
            break;
        }
        case 0x73:
        {
            execute_system(hart, instruction);
            hart->pc += 4;
            break;
        }
        default:
            assert(!"Unhandled opcode");
    }

    assert(hart->regs[0] == 0);
}

void execute_from(Hart* hart, uint32_t start_address)
{
    hart->pc = start_address;
    hart->halted = false;

    while (!hart->halted)
    {
        uint32_t instruction = load_word(hart, hart->pc);
        execute(hart, instruction);
    }
}