mod core;
mod mem;
mod state;

use wasm_bindgen::prelude::wasm_bindgen;
use crate::core::*;
use crate::mem::*;
use crate::state::CpuState;

const REGISTER_NAME_WIDTH: u32 = 0b11111;

const DRAM_LOAD: u32 = 0x20000000;
const OS_LOAD: u32 = 0x80000000;

const CONSOLE_LOAD: u32 = 0x70000000;

const INITRD_LOAD: u32 = 0;

#[wasm_bindgen(module = "/src/dev.js")]
extern "C" {
    fn uart_push_line(line: &str);
}

pub struct Uart {
    pub addr: u32,
    pub value: u8,
    acc: String
}

impl MemoryDevice for Uart {
    fn read(&self, address: u32) -> Result<u8, ()> {
        if address == self.addr {
            Ok(self.value)
        } else {
            Err(())
        }
    }

    fn write(&mut self, address: u32, value: u8) -> Result<(), ()> {
        if address == self.addr {
            self.value = value;

            if self.value as char == '\n' {
                uart_push_line(self.acc.as_str());
                self.acc = String::new();
            } else {
                self.acc.push(self.value as char);
            }
            Ok(())
        } else {
            Err(())
        }
    }
}

#[wasm_bindgen]
pub fn boot(image: &[u8], initrd: &[u8], memory: &[u8]) {
    let mut cpu_memory = CpuMemory {
        devices: Vec::new(),
    };

    let memory_len = memory.len();
    let os_size = image.len();
    let initrd_size = initrd.len();

    let ram_device = LogicalMemory {
        base: DRAM_LOAD,
        size: memory_len as u32,
        dram: memory.to_vec().into_boxed_slice(),
    };

    let os_image = LogicalMemory {
        base: OS_LOAD,
        size: os_size as u32,
        dram: image.to_vec().into_boxed_slice(),
    };

    let initrd_image = LogicalMemory {
        base: INITRD_LOAD,
        size: initrd_size as u32,
        dram: initrd.to_vec().into_boxed_slice(),
    };

    let uart = Uart {
        addr: CONSOLE_LOAD,
        value: 0,
        acc: "".to_string(),
    };

    cpu_memory.register_device(
        DRAM_LOAD..DRAM_LOAD + memory_len as u32,
        Box::new(ram_device)
    );
    cpu_memory.register_device(
        OS_LOAD..OS_LOAD + os_size as u32,
        Box::new(os_image)
    );
    cpu_memory.register_device(
        INITRD_LOAD..INITRD_LOAD + initrd_size as u32,
        Box::new(initrd_image)
    );
    cpu_memory.register_device(
        CONSOLE_LOAD..CONSOLE_LOAD + 1,
        Box::new(uart)
    );

    let mut state = CpuState::new(cpu_memory);

    state.pc = 0x80000000;
    state.registers[10] = 0;
    state.registers[2] = DRAM_LOAD + memory_len as u32;

    while state.on {
        {
            handle_instruction(&mut state);
            state.registers[0] = 0;
        }
    }
}

fn handle_instruction(state: &mut CpuState) {
    let pc = state.pc;
    let word = state.memory.read_instruction(pc).expect(format!("pc {:x} does not map to valid memory", pc).as_str());
    let opcode = word & 0x7f;
    match opcode {
        0b0110011 => {
            let f7 = word >> 25;
            let f3 = (word >> 12) & 0b111;

            let rs2 = ((word >> 20) & REGISTER_NAME_WIDTH) as usize;
            let rs1 = ((word >> 15) & REGISTER_NAME_WIDTH) as usize;
            let rd = ((word >> 7) & REGISTER_NAME_WIDTH) as usize;

            match (f7 << 3) | f3 {
                0 => add(state, rd, rs1, rs2),
                256 => sub(state, rd, rs1, rs2),
                4 => xor(state, rd, rs1, rs2),
                6 => or(state, rd, rs1, rs2),
                7 => and(state, rd, rs1, rs2),
                1 => sll(state, rd, rs1, rs2),
                5 => srl(state, rd, rs1, rs2),
                261 => sra(state, rd, rs1, rs2),
                2 => slt(state, rd, rs1, rs2),
                3 => sltu(state, rd, rs1, rs2),
                8 => mul(state, rd, rs1, rs2),
                9 => mulh(state, rd, rs1, rs2),
                10 => mulhsu(state, rd, rs1, rs2),
                11 => mulhu(state, rd, rs1, rs2),
                12 => div(state, rd, rs1, rs2),
                13 => divu(state, rd, rs1, rs2),
                14 => rem(state, rd, rs1, rs2),
                15 => remu(state, rd, rs1, rs2),
                _ => panic!("Illegal instruction: {:x}", word),
            }
        }
        0b0010011 => {
            let imm = ((word as i32) >> 20) as u32;
            let f3 = (word >> 12) & 0b111;

            let rs1 = ((word >> 15) & 0b11111) as usize;
            let rd = ((word >> 7) & 0b11111) as usize;

            match f3 {
                0 => addi(state, rd, rs1, imm),
                4 => xori(state, rd, rs1, imm),
                6 => ori(state, rd, rs1, imm),
                7 => andi(state, rd, rs1, imm),
                1 => {
                    if (imm >> 5) != 0 {
                        panic!("Illegal instruction: {:x}", word);
                    }
                    slli(state, rd, rs1, imm);
                }
                5 => {
                    if (imm >> 5) != 0 {
                        panic!("Illegal instruction: {:x}", word);
                    }
                    if (imm >> 5) & 0b1111111 == 0 {
                        srli(state, rd, rs1, imm);
                    } else {
                        srai(state, rd, rs1, imm);
                    }
                }
                2 => slti(state, rd, rs1, imm),
                3 => sltiu(state, rd, rs1, imm),
                _ => panic!("Illegal instruction: {:x}", word),
            }
        }
        0b0000011 => {
            let imm = ((word as i32) >> 20) as u32;
            let f3 = (word >> 12) & 0b111;

            let rs1 = ((word >> 15) & 0b11111) as usize;
            let rd = ((word >> 7) & 0b11111) as usize;

            match f3 {
                0 => lb(state, rd, rs1, imm),
                1 => lh(state, rd, rs1, imm),
                2 => lw(state, rd, rs1, imm),
                4 => lbu(state, rd, rs1, imm),
                5 => lhu(state, rd, rs1, imm),
                _ => panic!("Illegal instruction: {:x}", word),
            }
        }
        0b0100011 => {
            let imm_low = (word >> 7) & 0b11111;
            let imm_high = (word >> 25) & 0b1111111;
            let f3 = (word >> 12) & 0b111;
            let rs1 = ((word >> 15) & REGISTER_NAME_WIDTH) as usize;
            let rs2 = ((word >> 20) & REGISTER_NAME_WIDTH) as usize;

            let imm = ((((imm_high << 5) | imm_low) as i32) << 20 >> 20) as u32;

            match f3 {
                0 => sb(state, rs1, rs2, imm),
                1 => sh(state, rs1, rs2, imm),
                2 => sw(state, rs1, rs2, imm),
                _ => panic!("Illegal instruction: {:x}", word),
            };
        }
        0b1100011 => {
            let f3 = (word >> 12) & 0b111;
            let rs1 = ((word >> 15) & REGISTER_NAME_WIDTH) as usize;
            let rs2 = ((word >> 20) & REGISTER_NAME_WIDTH) as usize;

            let imm = (((word >> 31) & 0x1) << 12)
                | (((word >> 7) & 0x1) << 11)
                | (((word >> 25) & 0x3f) << 5)
                | (((word >> 8) & 0xf) << 1);

            let imm = ((imm as i32) << 19 >> 19) as u32;

            match f3 {
                0 => beq(state, rs1, rs2, imm),
                1 => bne(state, rs1, rs2, imm),
                4 => blt(state, rs1, rs2, imm),
                5 => bge(state, rs1, rs2, imm),
                6 => bltu(state, rs1, rs2, imm),
                7 => bgeu(state, rs1, rs2, imm),
                _ => panic!("Illegal instruction: {:x}", word),
            }
            return;
        }
        0b1101111 => {
            let rd = (word >> 7) & REGISTER_NAME_WIDTH;
            let imm = (((word >> 31) & 0x1) << 20)
                | (((word >> 12) & 0xFF) << 12)
                | (((word >> 20) & 0x1) << 11)
                | (((word >> 21) & 0x3FF) << 1);
            let imm = ((imm as i32) << 11 >> 11) as u32;
            jal(state, rd as usize, imm);
            return;
        }
        0b1100111 => {
            let rs1 = (word >> 15) & REGISTER_NAME_WIDTH;
            let rd = (word >> 7) & REGISTER_NAME_WIDTH;
            let imm = word >> 20;

            let sentinel = (word >> 12) & 0b111;
            if sentinel != 0 {
                panic!("Illegal instruction: {:x}", word);
            }

            jalr(state, rd as usize, rs1 as usize, imm);
            return;
        }
        0b0110111 => {
            let imm = word >> 12;
            let rd = ((word >> 7) & 0b11111) as usize;
            lui(state, rd, imm);
        }
        0b0010111 => {
            let imm = word >> 12;
            let rd = (word >> 7) & REGISTER_NAME_WIDTH;
            auipc(state, rd as usize, imm);
        }
        0b1110011 => {
            let imm = word >> 20;
            let rd = ((word >> 7) & REGISTER_NAME_WIDTH) as usize;
            let rs1 = ((word >> 15) & REGISTER_NAME_WIDTH) as usize;
            let f3 = (word >> 12) & 0b111;

            match f3 {
                0 => {
                    match imm {
                        0 => trap_to_kernel(state),        // ecall
                        258 => return_to_userspace(state), // sret
                        261 => {
                            state.on = false;
                            return;
                        }
                        _ => panic!("Illegal instruction: {:x}", word),
                    }
                } // ecall
                1 => csrrw(state, rd, rs1, imm as usize),
                2 => csrrs(state, rd, rs1, imm as usize),
                3 => csrrc(state, rd, rs1, imm as usize),
                5 => csrrwi(state, rd, rs1 as u32, imm as usize),
                6 => csrrsi(state, rd, rs1 as u32, imm as usize),
                7 => csrrci(state, rd, rs1 as u32, imm as usize),

                _ => panic!("Illegal instruction: {:x}", word),
            }
        } // ecall (ebreak not implemented)
        _ => panic!("Illegal instruction or todo: {:032b}", word),
    }

    state.pc += 4;
}

fn trap_to_kernel(state: &mut CpuState) {
    state.csr[0x341 /* mepc */] = state.pc;
    state.pc = state.csr[0x305 /* mtvec */];
}

fn return_to_userspace(state: &mut CpuState) {
    state.pc = state.csr[0x341 /* mepc */] + 4;
}