use anyhow::{anyhow, Result};

pub struct RV32 {
    regs: [u32; 32], // registers
    mem: Vec<u8>,    // memory
    pc: u32,
}

impl RV32 {
    pub fn new(mem: Vec<u8>) -> Self {
        Self {
            regs: [0; 32],
            pc: 0,
            mem,
        }
    }
    pub fn set_register(&mut self, index: usize, value: u32) {
        assert!(index < 32, "register index must <32");
        if index != 0 {
            self.regs[index] = value;
        }
    }
    // validate the access
    fn range(&self, address: u32, size: usize) -> Result<std::ops::Range<usize>> {
        if !address.is_multiple_of(size as u32) {
            return Err(anyhow!("misaligned at address: {}", address));
        }
        let start = address as usize;
        let end = start
            .checked_add(size)
            .filter(|&end| end <= self.mem.len())
            .ok_or(anyhow!("access fault, address: {}", address))?;
        Ok(start..end)
    }
    fn read(&self, address: u32, size: usize) -> Result<u32> {
        let range = self.range(address, size)?;
        let mut bytes = [0; 4];
        bytes[..size].copy_from_slice(&self.mem[range]);
        Ok(u32::from_le_bytes(bytes))
    }
    pub fn step(&mut self) -> Result<()> {
        let instruction = self.read(self.pc, 4)?;
        let next_pc = self.pc.wrapping_add(4);

        // R-type instruction:
        //                             R-type:
        // 31        25 24        20 19        15 14    12 11     7 6        0
        // ┌───────────┬────────────┬────────────┬────────┬────────┬─────────┐
        // │   funct7  │    rs2     │    rs1     │ funct3 │   rd   │ opcode  │
        // └───────────┴────────────┴────────────┴────────┴────────┴─────────┘
        let rd = ((instruction >> 7) & 31) as usize;
        let funct3 = (instruction >> 12) & 7;
        let rs1 = ((instruction >> 15) & 31) as usize;
        let rs2 = ((instruction >> 20) & 31) as usize;
        let funct7 = instruction >> 25;

        let a = self.regs[rs1];
        let b = self.regs[rs2];
        let result: Option<u32>;

        // 0x7f=0111_1111, match the lowest 7 bits to read the opcode
        match instruction & 0x7f {
            // R-type (register-register) integer operations
            0x33 => {
                result = Some(match (funct7, funct3) {
                    // ADD: add rs1 and rs2, wrapping to 32 bits
                    (0, 0) => a.wrapping_add(b),
                    // SUB: subtract rs2 from rs1, wrapping to 32 bits
                    (0x20, 0) => a.wrapping_sub(b),
                    // SLL: shift rs1 left by the low five bits of rs2
                    (0, 1) => a << (b & 31),
                    // SLT: set rd to one if rs1 is less than rs2 as signed
                    // values, else zero
                    (0, 2) => u32::from((a as i32) < (b as i32)),
                    // SLTU: set rd to one if rs1 is less than rs2 as unsigned
                    // values, else zero
                    (0, 3) => u32::from(a < b),
                    // XOR: write the bitwise XOR of rs1 and rs2 into rd
                    (0, 4) => a ^ b,
                    // SRL: shift rs1 right by the low five bits of rs2, filling
                    // with zeros
                    (0, 5) => a >> (b & 31),
                    // SRA: shift rs1 right by the low five bits of rs2,
                    // preserving its sign
                    (0x20, 5) => ((a as i32) >> (b & 31)) as u32,
                    // OR: write the bitwise OR of rs1 and rs2 into rd
                    (0, 6) => a | b,
                    // AND: write the bitwise AND of rs1 and rs2 into rd
                    (0, 7) => a & b,
                    _ => return Err(anyhow!("illegal instruction (A): {}", instruction)),
                });
            }
            _ => return Err(anyhow!("illegal instruction (B): {}", instruction)),
        }

        if next_pc & 3 != 0 {
            return Err(anyhow!("misaligned, address: {}", next_pc,));
        }
        if let Some(value) = result {
            self.set_register(rd, value);
        }
        self.pc = next_pc;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    const OPCODE_R: u32 = 0b0110011;

    // encodes R-type instruction
    fn encode_r(name: &str, rd: u32, rs1: u32, rs2: u32) -> u32 {
        let (funct7, funct3) = match name {
            "add" => (0b0000000, 0b000),
            "sub" => (0b0100000, 0b000),
            "sll" => (0b0000000, 0b001),
            "slt" => (0b0000000, 0b010),
            "sltu" => (0b0000000, 0b011),
            "xor" => (0b0000000, 0b100),
            "srl" => (0b0000000, 0b101),
            "sra" => (0b0100000, 0b101),
            "or" => (0b0000000, 0b110),
            "and" => (0b0000000, 0b111),
            _ => panic!("unknown R-type instruction: {name}"),
        };

        (funct7 << 25) | (rs2 << 20) | (rs1 << 15) | (funct3 << 12) | (rd << 7) | OPCODE_R
    }

    #[test]
    fn test_r_type_operations() {
        let test_vectors = [
            // (op, rs1, rs2, rd)
            ("add", 10, 20, 30),
            ("add", u32::MAX, 1, 0), // overflow
            ("sub", 30, 20, 10),
            ("sub", 0, 1, u32::MAX), // overflow
            ("sll", 1, 4, 16),
            ("sll", 1, 31, 0x8000_0000), // SLL maximum shift
            ("sll", 1, 32, 1),           // SLL shift wrapping
            ("slt", (-10i32) as u32, 20, 1),
            ("sltu", 10, 20, 1),
            ("slt", 0xffff_ffff, 0, 1), // Signed comparison: -1 < 0
            ("slt", 0xffff_ffff, 0, 1), // Signed vs unsigned distinction
            ("sltu", 0xffff_ffff, 0, 0),
            ("xor", 0b1010, 0b1100, 0b0110),
            ("srl", 0x8000_0000, 4, 0x0800_0000),
            ("srl", 0x8000_0000, 1, 0x4000_0000), // Logical shift should zero-fill
            ("sra", 0x8000_0000, 4, 0xf800_0000),
            ("sra", 0x8000_0000, 1, 0xc000_0000), // Arithmetic shift should sign-fill
            ("sra", 0xffff_ffff, 31, 0xffff_ffff), // Arithmetic shift of -1 stays -1
            ("or", 0b1010, 0b1100, 0b1110),
            ("and", 0b1010, 0b1100, 0b1000),
        ];

        for (name, rs1_value, rs2_value, expected) in test_vectors {
            let rd = 3;
            let rs1 = 1;
            let rs2 = 2;

            let instruction = encode_r(name, rd, rs1, rs2);

            let mut cpu = RV32::new(instruction.to_le_bytes().to_vec());

            cpu.regs[rs1 as usize] = rs1_value;
            cpu.regs[rs2 as usize] = rs2_value;

            cpu.step().unwrap();

            assert_eq!(
                cpu.regs[rd as usize], expected,
                "{name} failed:
                 rs1={rs1_value:#010x}, rs2={rs2_value:#010x}, 
                 expected={expected:#010x}, got={:#010x}",
                cpu.regs[rd as usize],
            );
        }
    }
}