use elf::ElfBytes;
use elf::endian::AnyEndian;
use std::io;
use std::io::Write;
use std::process::exit;
use std::sync::{LazyLock, Mutex};

const MEMORY_SIZE: usize = 64 * 1024 * 1024;
static MEMORY: LazyLock<Mutex<Box<[u8]>>> = LazyLock::new(|| {
    Mutex::new(vec![0u8; MEMORY_SIZE].into_boxed_slice())
});

fn main() {
    let path = std::path::PathBuf::from("hi");

    let f_dat = std::fs::read(path).expect("Unable to read file");

    let elf = ElfBytes::<AnyEndian>::minimal_parse(f_dat.as_slice()).expect("Unable to parse elf");

    let (sh, stab) = elf
        .section_headers_with_strtab()
        .expect("Unable to get section headers");

    let dup = f_dat.clone();

    // write memory
    {
        let mut memory = MEMORY.lock().expect("Unable to get memory");
        for i in 0..dup.len() {
            memory[i] = dup[i];
        }
    }

    let mut registers: [u32; 32] = [0; 32];
    let mut pc: u32 = 0;

    let str_tab = stab.expect("Cannot read strtab");
    if let Some(headers) = sh {
        for header in headers {
            // Search for .text
            if let Ok(sh_name) = str_tab.get(header.sh_name as usize)
                && sh_name == ".text"
            {
                pc = header.sh_addr as u32;
            }
        }
    }

    loop {
        {
            let mut memory = MEMORY.lock().expect("Unable to get memory");
            handle_instruction(&mut *memory, &mut registers, &mut pc);
        }
    }
}

fn handle_instruction(memory: &mut Box<[u8]>, registers: &mut [u32], pc: &mut u32) {
    let word = u32::from_le_bytes([
        memory[*pc as usize],
        memory[*pc as usize + 1],
        memory[*pc as usize + 2],
        memory[*pc as usize + 3],
    ]);
    let opcode = word & 0x7f;
    match opcode {
        0b0110011 => {
            let f7 = word >> 25;
            let f3 = (word >> 12) & 0b111;

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

            match (f7 << 3) | f3 {
                0 => registers[rd] = registers[rs1].wrapping_add(registers[rs2]), // add
                256 => registers[rd] = registers[rs1].wrapping_sub(registers[rs2]), // sub
                4 => registers[rd] = registers[rs1] ^ registers[rs2], // xor
                6 => registers[rd] = registers[rs1] | registers[rs2], // or
                7 => registers[rd] = registers[rs1] & registers[rs2], // and
                1 => registers[rd] = registers[rs1] << (registers[rs2] & 0b11111), // sll
                5 => registers[rd] = registers[rs1] >> (registers[rs2] & 0b11111), // srl
                261 => registers[rd] = ((registers[rs1] as i32) >> (registers[rs2] & 0b11111)) as u32, // sra
                2 => {
                    let rs1_signed: i32 = registers[rs1] as i32;
                    let rs2_signed: i32 = registers[rs2] as i32;
                    registers[rd] = if rs1_signed < rs2_signed { 1 } else { 0 }; // slt
                }
                3 => {
                    registers[rd] = if registers[rs1] < registers[rs2] {
                        1
                    } else {
                        0
                    }
                } // sltu
                8 => registers[rd] = registers[rs1].wrapping_mul(registers[rs2]), // mul
                9 => registers[rd] = (((registers[rs1] as i32 as i64) * (registers[rs2] as i32 as i64)) >> 32) as u32, // mulh
                10 => {
                    let rs1_signed = (registers[rs1] as i32) as i64;
                    let rs2_unsigned = registers[rs2] as i64;
                    registers[rd] = ((rs1_signed * rs2_unsigned) >> 32) as u32;
                } // mulhsu
                11 => registers[rd] = ((registers[rs1] as u64 * registers[rs2] as u64) >> 32) as u32, // mulhu
                12 => {
                    let a = registers[rs1] as i32;
                    let b = registers[rs2] as i32;
                    registers[rd] = if b == 0 {
                        u32::MAX
                    } else if a == i32::MIN && b == -1 {
                        a as u32
                    } else {
                        (a / b) as u32
                    };
                }, // div
                13 => registers[rd] = if registers[rs2] == 0 {
                    u32::MAX
                } else {
                    registers[rs1] / registers[rs2]
                }, // divu
                14 => {
                    let a = registers[rs1] as i32;
                    let b = registers[rs2] as i32;
                    registers[rd] = if b == 0 {
                        a as u32
                    } else if a == i32::MIN && b == -1 {
                        0
                    } else {
                        (a % b) as u32
                    }
                }, // rem
                15 => registers[rd] = if registers[rs2] == 0 {
                    registers[rs1]
                } else {
                    registers[rs1] % registers[rs2]
                }, //remu

                _ => 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 => registers[rd] = registers[rs1].wrapping_add(imm), // addi
                4 => registers[rd] = registers[rs1] ^ imm, // xori
                6 => registers[rd] = registers[rs1] | imm, // ori
                7 => registers[rd] = registers[rs1] & imm, // andi
                1 => {
                    if (imm >> 5) != 0 {
                        panic!("Illegal instruction: {:x}", word);
                    }
                    registers[rd] = registers[rs1] << (imm & 0b11111); // slli
                }
                5 => {
                    if (imm >> 5) != 0 {
                        panic!("Illegal instruction: {:x}", word);
                    }
                    if (imm >> 5) & 0b1111111 == 0 {
                        registers[rd] = registers[rs1] >> (imm & 0b11111); // srli
                    } else {
                        registers[rd] = ((registers[rs1] as i32 >> (imm & 0b11111))) as u32; // srai
                    }
                }
                2 => {
                    registers[rd] = if (registers[rs1] as i32) < (imm as i32) {
                        1
                    } else {
                        0
                    }
                } // slti
                3 => registers[rd] = if registers[rs1] < imm { 1 } else { 0 }, // sltiu
                _ => 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 => {
                    let byte: u8 = memory[(registers[rs1] + imm) as usize];
                    registers[rd] = (byte as i8) as u32; // lb
                }
                1 => {
                    let low = memory[(registers[rs1] + imm) as usize];
                    let high = memory[(registers[rs1] + imm + 1) as usize];
                    registers[rd] = (((high as u16) << 8 | low as u16) as i16) as u32; // lh
                }
                2 => {
                    let addr = (registers[rs1] + imm) as usize;
                    let slice = memory.get(addr..addr + 4).unwrap();
                    let number: u32 = u32::from_le_bytes(slice.try_into().unwrap());
                    registers[rd] = number; // lw
                }
                4 => registers[rd] = memory[(registers[rs1] + imm) as usize] as u32, // lbu
                5 => {
                    let low = memory[(registers[rs1] + imm) as usize];
                    let high = memory[(registers[rs1] + imm + 1) as usize];
                    registers[rd] = ((high as u16) << 8 | low as u16) as u32; // lhu
                }
                _ => 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) & 0b11111) as usize;
            let rs2 = ((word >> 20) & 0b11111) as usize;

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

            match f3 {
                0 => memory[(registers[rs1] + imm) as usize] = registers[rs2] as u8, // sb
                1 => {
                    let halfword = registers[rs2] & 0xFFFF;
                    let index = (registers[rs1] + imm) as usize;
                    memory[index] = (halfword & 0xFF) as u8;
                    memory[index + 1] = (halfword >> 8) as u8;
                } // sh
                2 => {
                    let index = (registers[rs1] + imm) as usize;
                    memory[index..index + 4].copy_from_slice(&registers[rs2].to_le_bytes()); // sw
                }
                _ => panic!("Illegal instruction: {:x}", word),
            };
        }
        0b1100011 => {
            let f3 = (word >> 12) & 0b111;
            let rs1 = ((word >> 15) & 0b11111) as usize;
            let rs2 = ((word >> 20) & 0b11111) 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 => {
                    if registers[rs1] == registers[rs2] {
                        *pc += imm;
                    }
                    return;
                } // beq
                1 => {
                    if registers[rs1] != registers[rs2] {
                        *pc += imm;
                    }
                    return;
                } // bne
                4 => {
                    if (registers[rs1] as i32) < (registers[rs2] as i32) {
                        *pc += imm;
                    }
                    return;
                } // blt
                5 => {
                    if (registers[rs1] as i32) >= (registers[rs2] as i32) {
                        *pc += imm;
                    }
                    return;
                } // bge
                6 => {
                    if registers[rs1] < registers[rs2] {
                        *pc += imm;
                    }
                    return;
                } // bltu
                7 => {
                    if registers[rs1] >= registers[rs2] {
                        *pc += imm;
                    }
                    return;
                } // bgeu
                _ => panic!("Illegal instruction: {:x}", word),
            }
        }
        0b1101111 => {
            let rd = (word >> 7) & 0b11111;
            let imm = word >> 12;
            registers[rd as usize] = *pc + 4;
            *pc += imm;
        } // jal
        0b1100111 => {
            let rs1 = (word >> 15) & 0b11111;
            let rd = (word >> 7) & 0b11111;
            let imm = word >> 20;

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

            registers[rd as usize] = *pc + 4;
            *pc = registers[rs1 as usize] + imm;
        }
        0b0110111 => {
            let imm = word >> 12;
            let rd = ((word >> 7) & 0b11111) as usize;
            registers[rd] = imm << 12; // lui
        }
        0b0010111 => {
            let imm = word >> 12;
            let rd = (word >> 7) & 0b11111;
            registers[rd as usize] = *pc + (imm << 12); // auipc
        }
        0b1110011 => {
            // ecall
            eprintln!(
                "TRAP! syscall {} with gp registers {:?} calling instruction {:x}",
                registers[17],
                &registers[10..17],
                *pc
            );
            if registers[17] == 0 {
                // sys_write: just a debug
                let buf = registers[11] as usize;
                let len = registers[12] as usize;
                let total_bytes = &memory[buf..buf + len];
                io::stdout().write(&total_bytes).unwrap();
            } else if registers[17] == 60 {
                exit(registers[10] as i32);
            }
        } // ecall (ebreak not implemented)
        _ => panic!("Illegal instruction or todo: {:x}", word),
    }

    *pc += 4;
}