use elf::ElfBytes;
use elf::endian::AnyEndian;
use std::io;
use std::io::Write;
use std::path::Component::ParentDir;
use std::sync::{LazyLock, Mutex};

const ABI: [&str; 32] = [
    "x0", "ra", "sp", "gp", "tp", "t0", "t1", "t2", "s0", "s1", "a0", "a1", "a2", "a3", "a4", "a5",
    "a6", "a7", "s2", "s3", "s4", "s5", "s6", "s7", "s8", "s9", "s10", "s11", "t3", "t4", "t5",
    "t6",
];

static MEMORY: LazyLock<Mutex<[u8; 64 * 1024 * 1024]>> =
    LazyLock::new(|| Mutex::new([0u8; 64 * 1024 * 1024]));

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 [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] + registers[rs2], // add
                256 => registers[rd] = registers[rs1] - 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], // sll
                5 => registers[rd] = registers[rs1] >> registers[rs2], // srl
                261 => registers[rd] = ((registers[rs1] as i32) >> registers[rs2]) 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
                _ => 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] + 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 (word >> 5) != 0 {
                        panic!("Illegal instruction: {:x}", word);
                    }
                    registers[rd] = registers[rs1] << (imm & 0b11111); // slli
                }
                5 => {
                    if (word >> 5) != 0 {
                        panic!("Illegal instruction: {:x}", word);
                    }
                    if (imm >> 5) & 0b1111111 == 0 {
                        registers[rd] = registers[rs1] >> (imm & 0b11111); // srli
                    } else {
                        registers[rd] = ((registers[rs1] >> (imm & 0b11111)) as i32) 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 >> 20;
            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 = slice.iter().fold(0u32, |acc, &n| (acc << 8) | n as u32);
                    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;

            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 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;

            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 => println!("J-type (todo)"),
        0b1100111 => println!("jalr"),
        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();
            }
        }
        _ => panic!("Illegal instruction or todo: {:x}", word),
    }

    *pc += 4;
}