use crate::{runtime_error::{KabelRuntimeError, RuntimeErrorKind}, vm_boolean_comparison, vm_boolean_equality};
use crate::vm_error;

#[derive(Debug, Clone)]
pub struct Unit {
    pub code: Vec<u8>,
    pub pool: Vec<Value>,
    pub lines: Vec<(usize, usize)>, // line #, repeats number of instructions
}
impl Unit {
    pub fn new(code: Vec<u8>, pool: Vec<Value>, lines: Vec<(usize, usize)>) -> Self {
        Self {
            code,
            pool,
            lines,
        }
    }
    pub fn new_empty() -> Self {
        Self {
            code: Vec::new(),
            pool: Vec::new(),
            lines: Vec::new(),
        }
    }
}

#[derive(Clone)]
pub struct VM {
    pub ip: usize,
    pub unit_ptr: usize,
    pub stack_ptr: usize,
    pub call_stack: Vec<(usize, usize, usize)>, // (unit_ptr, ip, stack_offset)
    pub units: Vec<Unit>,
    pub stack: Vec<Value>,
    pub stack_offset: usize,
    text: Vec<String>
}

impl VM {
    pub fn new(bytecode: Vec<u8>, lines: Vec<(usize, usize)>, pool: Vec<Value>,
        text: String) -> Self {
        Self {
            ip: 0,
            unit_ptr: 0,
            stack_ptr: 0,
            call_stack: vec![(0, 0, 0)],
            units: vec![Unit::new(bytecode, pool, lines)],
            stack: Vec::new(),
            stack_offset: 0,
            text: text.lines().map(|s| s.to_string()).collect::<Vec<String>>(),
        }
    }
    pub fn run(&mut self, output: &mut String) -> Result<(), KabelRuntimeError> {
        use Value::*;
        while self.ip < self.units[self.unit_ptr].code.len() {
            match self.read() {
                0x00 => { // LOAD
                    let byte = self.read() as usize;
                    let value = self.units[self.unit_ptr].pool[byte].clone();
                    self.stack.push(value);
                }
                0x01 => { // VAR
                    let ptr = self.read() as usize;
                    let value = self.stack[ptr+self.stack_offset].clone();
                    //println!("offset {} ptr {}", offset, ptr);
                    //println!("stack {:?}", self.stack);
                    self.stack.push(value);
                }
                // ASSIGN
                0x02 => {
                    let value = self.stack.pop().unwrap();
                    let ptr = self.read();
                    let offset = self.call_stack.last().expect("var call stack last").2;
                    self.stack[ptr as usize + offset] = value.clone();
                    self.stack.push(value);
                }
                0x03 => { // ADD
                    match (self.stack.pop().unwrap(), self.stack.pop().unwrap()) {
                        (Num(v1), Num(v2)) => self.stack.push(Num(v1 + v2)),
                        (Str(v1), Str(v2)) => {
                            self.stack.push(Str(v1 + &v2));
                        },
                        (Bool(_v1), Bool(_v2)) => {
                            return Err(vm_error!(self, RuntimeErrorKind::WrongType, "Cannot add booleans"))
                        }
                        (v1, v2) => {
                            return Err(vm_error!(self, RuntimeErrorKind::MismatchedTypes, "Mismatched types: {} and {}", v1.type_str(), v2.type_str()))
                        }
                    }
                }
                0x04 => { // SUB
                    match (self.stack.pop().unwrap(), self.stack.pop().unwrap()) {
                        (Num(v1), Num(v2)) => self.stack.push(Num(v1 - v2)),
                        (Str(_v1), Str(_v2)) => {
                            return Err(vm_error!(self, RuntimeErrorKind::WrongType, "Cannot subtract strings"))
                        },
                        (Bool(_v1), Bool(_v2)) => {
                            return Err(vm_error!(self, RuntimeErrorKind::WrongType, "Cannot subtract booleans"))
                        }
                        (v1, v2) => {
                            return Err(vm_error!(self, RuntimeErrorKind::MismatchedTypes, "Mismatched types: {} and {}", v1.type_str(), v2.type_str()))
                        }
                    }
                }
                0x05 => { // MUL
                    match (self.stack.pop().unwrap(), self.stack.pop().unwrap()) {
                        (Num(v1), Num(v2)) => self.stack.push(Num(v1 * v2)),
                        (Str(v1), Num(v2)) => {
                            if v2.fract() == 0.0 {
                                self.stack.push(Str(v1.repeat(v2 as usize)));
                            } else {
                                return Err(vm_error!(self, RuntimeErrorKind::WrongType, "Number must be an integer"))
                            }
                        }
                        (Str(_v1), Str(_v2)) => {
                            return Err(vm_error!(self, RuntimeErrorKind::WrongType, "Cannot multiply strings"))
                        },
                        (Bool(_v1), Bool(_v2)) => {
                            return Err(vm_error!(self, RuntimeErrorKind::WrongType, "Cannot multiply booleans"))
                        }
                        (v1, v2) => {
                            return Err(vm_error!(self, RuntimeErrorKind::MismatchedTypes, "Mismatched types: {} and {}", v1.type_str(), v2.type_str()))
                        }
                    }
                }
                0x06 => { // DIV
                    match (self.stack.pop().unwrap(), self.stack.pop().unwrap()) {
                        (Num(v1), Num(v2)) => self.stack.push(Num(v1 / v2)),
                        (Str(_v1), Str(_v2)) => {
                            return Err(vm_error!(self, RuntimeErrorKind::WrongType, "Cannot divide strings"))
                        },
                        (Bool(_v1), Bool(_v2)) => {
                            return Err(vm_error!(self, RuntimeErrorKind::WrongType, "Cannot divide booleans"))
                        }
                        (v1, v2) => {
                            return Err(vm_error!(self, RuntimeErrorKind::MismatchedTypes, "Mismatched types: {} and {}", v1.type_str(), v2.type_str()))
                        }
                    }
                }
                0x07 => { // MOD
                    match (self.stack.pop().unwrap(), self.stack.pop().unwrap()) {
                        (Num(v1), Num(v2)) => self.stack.push(Num(v1 % v2)),
                        (Str(_v1), Str(_v2)) => {
                            return Err(vm_error!(self, RuntimeErrorKind::WrongType, "Cannot perform modulus on strings"))
                        },
                        (Bool(_v1), Bool(_v2)) => {
                            return Err(vm_error!(self, RuntimeErrorKind::WrongType, "Cannot perform modulus on booleans"))
                        }
                        (v1, v2) => {
                            return Err(vm_error!(self, RuntimeErrorKind::MismatchedTypes, "Mismatched types: {} and {}", v1.type_str(), v2.type_str()))
                        }
                    }
                }
                0x08 => { // BITAND
                    match (self.stack.pop().unwrap(), self.stack.pop().unwrap()) {
                        (Num(v1), Num(v2)) => {
                            if v1.fract() != 0.0 {
                                return Err(vm_error!(self, RuntimeErrorKind::WrongType, "Cannot perform bitwise AND on {}", v1))
                            }
                            if v2.fract() != 0.0 {
                                return Err(vm_error!(self, RuntimeErrorKind::WrongType, "Cannot perform bitwise AND on {}", v2))
                            }
                            self.stack.push(Num((v1 as u32 & v2 as u32) as f32))
                        }
                        (Str(_v1), Str(_v2)) => {
                            return Err(vm_error!(self, RuntimeErrorKind::WrongType, "Cannot perform bitwise AND on strings"))
                        },
                        (Bool(_v1), Bool(_v2)) => {
                            return Err(vm_error!(self, RuntimeErrorKind::WrongType, "Cannot perform bitwise AND on booleans"))
                        }
                        (v1, v2) => {
                            return Err(vm_error!(self, RuntimeErrorKind::MismatchedTypes, "Mismatched types: {} and {}", v1.type_str(), v2.type_str()))
                        }
                    }
                }
                0x09 => { // BITXOR
                    match (self.stack.pop().unwrap(), self.stack.pop().unwrap()) {
                        (Num(v1), Num(v2)) => {
                            if v1.fract() != 0.0 {
                                return Err(vm_error!(self, RuntimeErrorKind::WrongType, "Cannot perform bitwise XOR on {}", v1))
                            }
                            if v2.fract() != 0.0 {
                                return Err(vm_error!(self, RuntimeErrorKind::WrongType, "Cannot perform bitwise XOR on {}", v2))
                            }
                            self.stack.push(Num((v1 as u32 ^ v2 as u32) as f32))
                        }
                        (Str(_v1), Str(_v2)) => {
                            return Err(vm_error!(self, RuntimeErrorKind::WrongType, "Cannot perform bitwise XOR on strings"))
                        },
                        (Bool(_v1), Bool(_v2)) => {
                            return Err(vm_error!(self, RuntimeErrorKind::WrongType, "Cannot perform bitwise XOR on booleans"))
                        }
                        (v1, v2) => {
                            return Err(vm_error!(self, RuntimeErrorKind::MismatchedTypes, "Mismatched types: {} and {}", v1.type_str(), v2.type_str()))
                        }
                    }
                }
                0x0A => { // BITOR
                    match (self.stack.pop().unwrap(), self.stack.pop().unwrap()) {
                        (Num(v1), Num(v2)) => {
                            if v1.fract() != 0.0 {
                                return Err(vm_error!(self, RuntimeErrorKind::WrongType, "Cannot perform bitwise OR on {}", v1))
                            }
                            if v2.fract() != 0.0 {
                                return Err(vm_error!(self, RuntimeErrorKind::WrongType, "Cannot perform bitwise OR on {}", v2))
                            }
                            self.stack.push(Num((v1 as u32 | v2 as u32) as f32))
                        }
                        (Str(_v1), Str(_v2)) => {
                            return Err(vm_error!(self, RuntimeErrorKind::WrongType, "Cannot perform bitwise OR on strings"))
                        },
                        (Bool(_v1), Bool(_v2)) => {
                            return Err(vm_error!(self, RuntimeErrorKind::WrongType, "Cannot perform bitwise OR on booleans"))
                        }
                        (v1, v2) => {
                            return Err(vm_error!(self, RuntimeErrorKind::MismatchedTypes, "Mismatched types: {} and {}", v1.type_str(), v2.type_str()))
                        }
                    }
                }
                // EQ
                0x0b => vm_boolean_equality!(self, ==),
                // NE
                0x0C => vm_boolean_equality!(self, !=),
                // GR
                0x0D => vm_boolean_comparison!(self, >),
                // GE
                0x0E => vm_boolean_comparison!(self, >=),
                // LS
                0x0F => vm_boolean_comparison!(self, <),
                // LE
                0x10 => vm_boolean_comparison!(self, <=),
                // OR
                0x11 => match (self.stack.pop().unwrap(), self.stack.pop().unwrap()) {
                    (Num(_v1), Num(_v2)) => {
                        return Err(vm_error!(self, RuntimeErrorKind::WrongType, "Tried to perform \"||\" on non-boolean numbers"))
                    }
                    (Str(_v1), Str(_v2)) => {
                        return Err(vm_error!(self, RuntimeErrorKind::WrongType, "Tried to perform \"||\" on non-boolean strings"))
                    },
                    (Bool(v1), Bool(v2)) => self.stack.push(Bool(v1 || v2)),
                    (Null, _) => return Err(vm_error!(self, RuntimeErrorKind::WrongType, "Tried to perform \"||\" on non-boolean value null")),
                    (_, Null) => return Err(vm_error!(self, RuntimeErrorKind::WrongType, "Tried to perform \"||\" on non-boolean value null")),
                    (v1, v2) => {
                        return Err(vm_error!(self, RuntimeErrorKind::MismatchedTypes, "Mismatched types: {} and {}", v1.type_str(), v2.type_str()))
                    }
                }
                // AND
                0x12 => match (self.stack.pop().unwrap(), self.stack.pop().unwrap()) {
                    (Num(_v1), Num(_v2)) => {
                        return Err(vm_error!(self, RuntimeErrorKind::WrongType, "Tried to perform \"&&\" on non-boolean numbers"))
                    }
                    (Str(_v1), Str(_v2)) => {
                        return Err(vm_error!(self, RuntimeErrorKind::WrongType, "Tried to perform \"&&\" on non-boolean strings"))
                    },
                    (Bool(v1), Bool(v2)) => self.stack.push(Bool(v1 && v2)),
                    (Null, _) => return Err(vm_error!(self, RuntimeErrorKind::WrongType, "Tried to perform \"&&\" on non-boolean value null")),
                    (_, Null) => return Err(vm_error!(self, RuntimeErrorKind::WrongType, "Tried to perform \"&&\" on non-boolean value null")),
                    (v1, v2) => {
                        return Err(vm_error!(self, RuntimeErrorKind::MismatchedTypes, "Mismatched types: {} and {}", v1.type_str(), v2.type_str()))
                    }
                }
                // NOT
                0x13 => match self.stack.pop().unwrap() {
                    Null => { return Err(vm_error!(self, RuntimeErrorKind::WrongType, "Tried to perform \"!\" on non-boolean value null")) }
                    Num(_v1) => {
                        return Err(vm_error!(self, RuntimeErrorKind::WrongType, "Tried to perform \"!\" on non-boolean number"))
                    }
                    Str(_v1) => {
                        return Err(vm_error!(self, RuntimeErrorKind::WrongType, "Tried to perform \"!\" on non-boolean string"))
                    }
                    Bool(v1) => self.stack.push(Bool(!v1)),
                    Fun(_v1) => {
                        return Err(vm_error!(self, RuntimeErrorKind::WrongType, "Tried to perform \"!\" on non-boolean function"))
                    }
                }
                // NEG
                0x14 => match self.stack.pop().unwrap() {
                    Null => { return Err(vm_error!(self, RuntimeErrorKind::WrongType, "Tried to perform \"-\" on non-number value null")) }
                    Num(v1) => self.stack.push(Num(-v1)),
                    Str(_v1) => {
                        return Err(vm_error!(self, RuntimeErrorKind::WrongType, "Tried to perform \"-\" on non-number string"))
                    }
                    Bool(_v1) => {
                        return Err(vm_error!(self, RuntimeErrorKind::WrongType, "Tried to perform \"-\" on non-number boolean"))
                    }
                    Fun(_v1) => {
                        return Err(vm_error!(self, RuntimeErrorKind::WrongType, "Tried to perform \"-\" on non-number function"))
                    }
                }
                // JMP
                0x15 => {
                    let loc = self.read_u16();
                    self.ip += loc as usize;
                }
                // JMP_UP
                0x16 => {
                    let loc = self.read_u16();
                    self.ip -= loc as usize;
                }
                // JNE
                0x17 => {
                    let condition = self.stack.pop().unwrap();
                    if let Value::Bool(condition) = condition {
                        if !condition {
                            let loc = self.read_u16();
                            self.ip += loc as usize;
                        } else {
                            self.read_u16();
                        }
                    } else {
                        return Err(vm_error!(self, RuntimeErrorKind::WrongType, "if must have condition of type boolean"))
                    }
                }

                // CALL
                0x18 => {
                    let num_args = self.read();
                    let function = self.stack.pop().expect("Stack was empty in call");
                    if let Value::Fun(function) = function {
                        if num_args as usize != function.arity {
                            return Err(vm_error!(self, RuntimeErrorKind::IncorrectArity, "Function has {} arguments, {} provided", function.arity, num_args))
                        }
                        self.stack_offset = self.stack.len()-1;
                        self.call_stack.push((self.unit_ptr, self.ip, self.stack_offset));
                        for _ in 0..num_args {
                            let arg = self.stack.pop().expect("Missing arguments in call");
                            self.stack.push(arg);
                        }
                        self.stack.insert(self.stack.len()-num_args as usize, Value::Fun(function));
                        //println!("stack call {:?}", self.stack);
                        self.stack_ptr += 1;
                        self.ip = 0;
                        self.unit_ptr = function.unit_ptr as usize;
                    } else {
                        return Err(vm_error!(self, RuntimeErrorKind::WrongType, "Tried to call non-function type {}", function.type_str()))
                    }
                }
                // RET
                0x19 => {
                    let (unit_ptr, ip, stack_offset) = self.call_stack.pop().expect("Call stack empty on RET");
                    let ret = self.stack.pop().expect("Missing return value");
                    self.stack = self.stack[..stack_offset].to_vec();
                    self.stack_offset = self.call_stack.last().expect("call stack empty").2;
                    self.unit_ptr = unit_ptr;
                    self.ip = ip;
                    // returned to code
                    self.stack_ptr -= 1;
                    self.stack.push(ret);
                }

                0xFD => { // POP
                    let times = self.read();
                    for _ in 0..times {
                        self.stack.pop().expect(&format!("{}: Unable to pop stack", self.ip));
                    }
                }
                0xFE => { // PRINT
                    let value = self.stack.pop().unwrap();
                    match value {
                        Null => *output += "null",
                        Num(v) => *output += &v.to_string(),
                        Str(v) => *output += &v,
                        Bool(v) => *output += &v.to_string(),
                        Fun(v) => *output += &format!("<function {}>", v.unit_ptr),
                    }
                    *output += "\n";
                }
                _ => {}
            }
        }
        Ok(())
    }
    pub fn read(&mut self) -> u8 {
        let byte = self.units[self.unit_ptr].code[self.ip];
        self.ip += 1;
        byte
    }
    pub fn read_u16(&mut self) -> u16 {
        let byte_one = (self.units[self.unit_ptr].code[self.ip] as u16) << 0x08;
        self.ip += 1;
        let byte_two = self.units[self.unit_ptr].code[self.ip] as u16;
        self.ip += 1;
        byte_one | byte_two
    }
    pub fn find_line(&self) -> usize { // returns line # at ip
        let mut line_ip = 0;
        for (line, rep) in self.units[self.unit_ptr].lines.clone() {
            if line_ip + rep > self.ip {
                return line;
            }
            line_ip += rep;
        }
        println!("{} {}", line_ip, self.ip);
        panic!("Something went wrong in finding line for error")
    }
}

#[derive(Debug, Clone)]
pub enum Value {
    Null, Num(f32), Str(String), Bool(bool), Fun(Function)
}

impl Value {
    pub fn type_str(&self) -> String {
        match self {
            Value::Null => "null".to_string(),
            Value::Num(_) => "number".to_string(),
            Value::Str(_) => "string".to_string(),
            Value::Bool(_) => "bool".to_string(),
            Value::Fun(_) => "function".to_string(),
        }
    }
}

#[derive(Debug, Clone, Copy)]
pub struct Function {
    pub unit_ptr: usize,
    pub arity: usize,
}