use std::io;
use std::io::Write;
use elf::endian::AnyEndian;
use elf::ElfBytes;

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",
];

fn main() {
    let path = std::path::PathBuf::from("/home/theavgeekbee/projects/webos/sample/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 memory: &mut [u8;0xFFFF] = &mut [0; 0xFFFF];
    let dup = f_dat.clone();
    for i in 0..dup.len() {
        memory[i] = dup[i];
    }

    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"
            {
                let (sh_data, _) = elf
                    .section_data(&header)
                    .expect("Unable to get section data");

                handle_instruction(sh_data, header.sh_addr as u32, memory);
            }
        }
    }
}

fn handle_instruction(sh_data: &[u8], instruction: u32, memory: &[u8;0xFFFF]) {
    let mut registers: [u32;32] = [0;32];
    let mut pc = instruction;

    for i in (0..sh_data.len()).step_by(4) {
        let word = u32::from_le_bytes([sh_data[i], sh_data[i + 1], sh_data[i + 2], sh_data[i + 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 => {
                        println!(
                            "add {} {} {}",
                            ABI[rd], ABI[rs1], ABI[rs2]
                        );
                        registers[rd] = registers[rs1] + registers[rs2];
                    },
                    256 => {
                        println!(
                            "sub {} {} {}",
                            ABI[rd], ABI[rs1], ABI[rs2]
                        );
                        registers[rd] = registers[rs1] - registers[rs2];
                    },
                    4 => {
                        println!(
                            "xor {} {} {}",
                            ABI[rd], ABI[rs1], ABI[rs2]
                        );
                        registers[rd] = registers[rs1] ^ registers[rs2];
                    },
                    6 => {
                        println!(
                            "or {} {} {}",
                            ABI[rd], ABI[rs1], ABI[rs2]
                        );
                        registers[rd] = registers[rs1] | registers[rs2];
                    },
                    7 => {
                        println!(
                            "and {} {} {}",
                            ABI[rd], ABI[rs1], ABI[rs2]
                        );
                        registers[rd] = registers[rs1] & registers[rs2];
                    },
                    1 => {
                        println!(
                            "sll {} {} {}",
                            ABI[rd], ABI[rs1], ABI[rs2]
                        );
                        registers[rd] = registers[rs1] << registers[rs2];
                    },
                    5 => {
                        println!(
                            "srl {} {} {}",
                            ABI[rd], ABI[rs1], ABI[rs2]
                        );
                        registers[rd] = registers[rs1] >> registers[rs2];
                    },
                    261 => {
                        println!(
                            "sra {} {} {}",
                            ABI[rd], ABI[rs1], ABI[rs2]
                        );
                        registers[rd] = ((registers[rs1] as i32) >> registers[rs2]) as u32;
                    },
                    2 => {
                        println!(
                            "slt {} {} {}",
                            ABI[rd], ABI[rs1], ABI[rs2]
                        );
                        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};
                    },
                    3 => {
                        println!(
                            "sltu {} {} {}",
                            ABI[rd], ABI[rs1], ABI[rs2]
                        );
                        registers[rd] = if registers[rs1] < registers[rs2] {1} else {0};
                    },
                    _ => println!("Invalid opcode"),
                }
            }
            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 => {
                        println!("addi {} {} {}", ABI[rd], ABI[rs1], imm);
                        registers[rd] = registers[rs1] + imm;
                    },
                    4 => {
                        println!("xori {} {} {}", ABI[rd], ABI[rs1], imm);
                        registers[rd] = registers[rs1] ^ imm;
                    },
                    6 => {
                        println!("ori {} {} {}", ABI[rd], ABI[rs1], imm);
                        registers[rd] = registers[rs1] | imm;
                    },
                    7 => {
                        println!("andi {} {} {}", ABI[rd], ABI[rs1], imm);
                        registers[rd] = registers[rs1] & imm;
                    },
                    1 => {
                        if (word >> 5) != 0 {
                            todo!("Illegal instruction (trap)");
                        }
                        println!(
                            "slli {} {} {}",
                            ABI[rd],
                            ABI[rs1],
                            imm & 0b11111
                        );
                        registers[rd] = registers[rs1] << (imm & 0b11111);
                    },
                    5 => {
                        if (word >> 5) != 0 {
                            todo!("Illegal instruction (trap)");
                        }
                        if (imm >> 5) & 0b1111111 == 0 {
                            println!(
                                "srli {} {} {}",
                                ABI[rd],
                                ABI[rs1],
                                imm & 0b11111
                            );
                            registers[rd] = registers[rs1] >> (imm & 0b11111);
                        } else {
                            println!(
                                "srai {} {} {}",
                                ABI[rd],
                                ABI[rs1],
                                imm & 0b11111
                            );
                            registers[rd] = ((registers[rs1]  >> (imm & 0b11111)) as i32) as u32;

                        }
                    }
                    2 => {
                        println!("slti {} {} {}", ABI[rd], ABI[rs1], imm);
                        registers[rd] = if (registers[rs1] as i32) < (imm as i32) {1} else {0};
                    },
                    3 => {
                        println!("sltiu {} {} {}", ABI[rd], ABI[rs1], imm);
                        registers[rd] = if registers[rs1] < imm {1} else {0};
                    },
                    _ => println!("Invalid opcode"),
                }
            }
            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 => {
                        println!(
                            "lb {} {} {}",
                            ABI[rd],
                            ABI[rs1],
                            imm & 0xFF
                        );
                        let byte: u8 = memory[(registers[rs1] + imm) as usize];
                        registers[rd] = (byte as i8) as u32;

                    },
                    1 => {
                        println!(
                            "lh {} {} {}",
                            ABI[rd],
                            ABI[rs1],
                            imm & 0xFFFF
                        );
                        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;
                    },
                    2 => {
                        println!(
                            "lw {} {} {}",
                            ABI[rd],
                            ABI[rs1],
                            imm & 0xFFFFFF
                        );
                        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;
                    },
                    4 => {
                        println!(
                            "lbu {} {} {}",
                            ABI[rd],
                            ABI[rs1],
                            imm & 0xFF
                        );

                    },
                    5 => println!(
                        "lhu {} {} {}",
                        ABI[rd],
                        ABI[rs1],
                        imm & 0xFFFF
                    ),
                    _ => println!("Invalid opcode"),
                }
            }
            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 => println!("sb {} {} {}", ABI[rs1], ABI[rs2], imm & 0xFF),
                    1 => println!("sh {} {} {}", ABI[rs1], ABI[rs2], imm & 0xFFFF),
                    2 => println!("sw {} {} {}", ABI[rs1], ABI[rs2], imm & 0xFFFFFF),
                    _ => println!("Invalid opcode")
                };

                // todo: in real kernel use memory pages
            },
            0b1100011 => println!("B-type"),
            0b1101111 => println!("J-type"),
            0b1100111 => println!("jalr"),
            0b0110111 => {
                let imm = word >> 12;
                let rd = ((word >> 7) & 0b11111) as usize;
                println!("lui {} {}", rd, imm as i32);
                registers[rd] = imm << 12;
            },
            0b0010111 => {
                let imm = word >> 12;
                let rd = (word >> 7) & 0b11111;
                println!("auipc {} {}", ABI[rd as usize], imm);
                registers[rd as usize] = pc + (imm << 12);
            }
            0b1110011 => {
                println!("ecall");
                eprintln!("TRAP! syscall {} with gp registers {:?} calling instruction {:x}", registers[17], &registers[10..17], pc);
                if registers[17] == 0 { // sys_write
                    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();
                }
            },
            _ => println!("Invalid opcode"),
        }

        pc += 4;
    }
}