use crate::{Machine, MEMORY};
use std::collections::HashMap;
use std::fs::read_to_string;

pub struct ImageAssembler {
    memory: [u32; MEMORY],
    ptr: usize,
    labels: HashMap<String, usize>,
}

impl ImageAssembler {
    pub fn new() -> ImageAssembler {
        ImageAssembler {
            memory: [0; MEMORY],
            ptr: 0,
            labels: HashMap::new(),
        }
    }

    pub fn extract(self) -> Machine {
        let mut machine = Machine::new();
        machine.memory = self.memory;
        return machine;
    }

    pub fn assemble(&mut self, filename: &str) {
        let asm = read_to_string(filename).unwrap_or_default();
        let lines = asm.split('\n');
        for line in lines {
            let parts = line.split_once(';');

            let command = match parts {
                Some((command, _comment)) => command.trim(),
                _ => line.trim(),
            };

            let command_parts = command.split_whitespace().collect::<Vec<&str>>();

            if command_parts.is_empty() {
                continue;
            }

            if command_parts.len() == 1 {
                match command_parts[0] {
                    "nop" => {
                        self.memory[self.ptr] = 0x00000000;
                        self.ptr += 1;
                        println!("{:X}: {:X}", self.ptr, self.memory[self.ptr]);
                    }
                    "ret" => {
                        self.memory[self.ptr] = 0xC0000000;
                        println!("{:X}: {:X}", self.ptr, self.memory[self.ptr]);
                        self.ptr += 1;
                    }
                    "hlt" => {
                        self.memory[self.ptr] = 0x0F000000;
                        self.ptr += 1;
                        println!("{:X}: {:X}", self.ptr, self.memory[self.ptr]);
                    }
                    "inc1" => {
                        self.memory[self.ptr] = 0x41000000;
                        self.ptr += 1;
                        println!("{:X}: {:X}", self.ptr, self.memory[self.ptr]);
                    }
                    "inc2" => {
                        self.memory[self.ptr] = 0x42000000;
                        self.ptr += 1;
                        println!("{:X}: {:X}", self.ptr, self.memory[self.ptr]);
                    }
                    "inc3" => {
                        self.memory[self.ptr] = 0x43000000;
                        self.ptr += 1;
                        println!("{:X}: {:X}", self.ptr, self.memory[self.ptr]);
                    }
                    "updsp" => {
                        self.memory[self.ptr] = 0x01000000;
                        self.ptr += 1;
                        println!("{:X}: {:X}", self.ptr, self.memory[self.ptr]);
                    }
                    _ => {
                        let data = self.parse_addr(command_parts[0]);
                        self.memory[self.ptr] = data;
                        self.ptr += 1;
                    }
                }
            } else if command_parts.len() == 2 && command_parts[0] != "data" {
                match command_parts[0] {
                    "label" => {
                        self.labels
                            .insert(command_parts[1].replace(":", ""), self.ptr);
                    }
                    "@" => {
                        self.ptr = self.parse_addr(command_parts[1]) as usize;
                        println!("Setting addr to ${:X}", self.ptr);
                    }
                    "inc" => {
                        let num = self.parse_addr(command_parts[1]);
                        let inst = 0x40 + num;
                        self.memory[self.ptr] = inst << 24;
                        println!("{:X}: {:X}", self.ptr, self.memory[self.ptr]);
                        self.ptr += 1;
                    }
                    "dec" => {
                        let num = self.parse_addr(command_parts[1]);
                        let inst = 0x50 + num;
                        self.memory[self.ptr] = inst << 24;
                        println!("{:X}: {:X}", self.ptr, self.memory[self.ptr]);
                        self.ptr += 1;
                    }
                    "jmp" => {
                        let addr = self.parse_addr(command_parts[1]);
                        self.memory[self.ptr] = 0xE0000000 + addr;
                        println!("{:X}: {:X}", self.ptr, self.memory[self.ptr]);
                        self.ptr += 1;
                    }
                    "jneq" => {
                        let addr = self.parse_addr(command_parts[1]);
                        self.memory[self.ptr] = 0xE2000000 + addr;
                        println!("{:X}: {:X}", self.ptr, self.memory[self.ptr]);
                        self.ptr += 1;
                    }
                    "call" => {
                        let addr = self.parse_addr(command_parts[1]);
                        self.memory[self.ptr] = 0xC1000000 + addr;
                        println!("{:X}: {:X}", self.ptr, self.memory[self.ptr]);
                        self.ptr += 1;
                    }
                    "import" => {
                        self.assemble(command_parts[1]);
                    }
                    _ => {
                        panic!("Unknown Instruction: {:?}", command_parts)
                    }
                }
            } else if command_parts.len() == 3 && command_parts[0] != "data" {
                match command_parts[0] {
                    "ldstore" => {
                        // Load Immediate
                        let target_reg = self.parse_addr(command_parts[1]);
                        let inst = (0x70 + target_reg) << 24;
                        let imm = self.parse_addr(command_parts[2]);
                        self.memory[self.ptr] = inst + imm;
                        println!("{:X}: {:X}", self.ptr, self.memory[self.ptr]);
                        self.ptr += 1;
                    }
                    "loadi" => {
                        // Load Immediate
                        let target_reg = self.parse_addr(command_parts[1]);
                        let inst = (0x30 + target_reg) << 24;
                        let imm = self.parse_addr(command_parts[2]);
                        self.memory[self.ptr] = inst + imm;
                        println!("{:X}: {:X}", self.ptr, self.memory[self.ptr]);
                        self.ptr += 1;
                    }
                    "load" => {
                        // Load Immediate
                        let target_reg = self.parse_addr(command_parts[1]);
                        let inst = (0x20 + target_reg) << 24;
                        let reg = self.parse_addr(command_parts[2]);
                        self.memory[self.ptr] = inst + reg;
                        println!("{:X}: {:X}", self.ptr, self.memory[self.ptr]);
                        self.ptr += 1;
                    }
                    "store" => {
                        // Load Immediate
                        let target_reg = self.parse_addr(command_parts[1]);
                        let inst = (0x10 + target_reg) << 24;
                        let reg = self.parse_addr(command_parts[2]);
                        self.memory[self.ptr] = inst + reg;
                        println!("{:X}: {:X}", self.ptr, self.memory[self.ptr]);
                        self.ptr += 1;
                    }
                    "rstore" => {
                        // Load Immediate
                        let target_reg = self.parse_addr(command_parts[1]);
                        let inst = (0x60 + target_reg) << 24;
                        let reg = self.parse_addr(command_parts[2]);
                        self.memory[self.ptr] = inst + reg;
                        println!("{:X}: {:X}", self.ptr, self.memory[self.ptr]);
                        self.ptr += 1;
                    }
                    "tx" => {
                        let target_reg = self.parse_addr(command_parts[1]);
                        let inst = (0x80 + target_reg) << 24;
                        let reg = self.parse_addr(command_parts[2]);
                        self.memory[self.ptr] = inst + reg;
                        println!("{:X}: {:X}", self.ptr, self.memory[self.ptr]);
                        self.ptr += 1;
                    }
                    _ => {
                        panic!("unknown opcpde {}", command_parts[0])
                    }
                }
            } else if command_parts.len() == 4 && command_parts[0] != "data" {
                // Every Instruction with 3 is [INST] [Source Reg] [Target Addr / Reg]
                let a_reg = self.parse_addr(command_parts[1]);
                let b_reg = self.parse_addr(command_parts[2]);
                let result_reg = self.parse_addr(command_parts[3]);
                match command_parts[0] {
                    "add" => {
                        self.memory[self.ptr] =
                            (0xa1 << 24) + (a_reg << 16) + (b_reg << 8) + result_reg;
                        println!("{:X}: {:X}", self.ptr, self.memory[self.ptr]);
                        self.ptr += 1;
                    }
                    "sub" => {
                        self.memory[self.ptr] =
                            (0xa2 << 24) + (a_reg << 16) + (b_reg << 8) + result_reg;
                        println!("{:X}: {:X}", self.ptr, self.memory[self.ptr]);
                        self.ptr += 1;
                    }
                    "mul" => {
                        self.memory[self.ptr] =
                            (0xa3 << 24) + (a_reg << 16) + (b_reg << 8) + result_reg;
                        println!("{:X}: {:X}", self.ptr, self.memory[self.ptr]);
                        self.ptr += 1;
                    }
                    "shift" => {
                        self.memory[self.ptr] =
                            (0xa5 << 24) + (a_reg << 16) + (b_reg << 8) + result_reg;
                        println!("{:X}: {:X}", self.ptr, self.memory[self.ptr]);
                        self.ptr += 1;
                    }
                    "and" => {
                        self.memory[self.ptr] =
                            (0xa6 << 24) + (a_reg << 16) + (b_reg << 8) + result_reg;
                        println!("{:X}: {:X}", self.ptr, self.memory[self.ptr]);
                        self.ptr += 1;
                    }
                    _ => {
                        panic!("Unknown Instruction: {:?}", command_parts)
                    }
                }
            } else {
                match command_parts[0] {
                    "data" => match command_parts[1] {
                        "ascii" => {
                            self.labels
                                .insert(command_parts[2].replace(":", ""), self.ptr);
                            let data = command_parts
                                .iter()
                                .skip(3)
                                .map(|x| String::from(*x))
                                .collect::<Vec<String>>()
                                .join(" ");
                            for char in data.chars() {
                                self.memory[self.ptr] = char as u32;
                                self.ptr += 1;
                            }
                            self.memory[self.ptr] = 0x00;
                            self.ptr += 1;
                        }
                        _ => {
                            panic!("unknown type {}", command_parts[1])
                        }
                    },
                    _ => {
                        panic!("unknown command {}", command_parts[0])
                    }
                }
            }
        }
    }

    fn parse_addr(&self, input: &str) -> u32 {
        match u32::from_str_radix(input.replace("$", "").replace(":", "").as_str(), 16) {
            Ok(addr) => {
                return addr;
            }
            _ => match self.labels.get(input) {
                Some(addr) => {
                    return addr.clone() as u32;
                }
                _ => {
                    panic!("Not a Hex Encoded Address {:?}", input)
                }
            },
        }
    }
}