pub mod compiled;
pub mod functions;

use std::collections::HashMap;

use super::errors::RuntimePanic;

pub use compiled::*;
pub use functions::*;

pub struct VirtualMachine {
    registry_stack: Vec<[Option<Value>; 8]>,
    registry: [Option<Value>; 8],

    stack: Vec<Value>,
    heap: HashMap<HeapID, AllocatedVar>,
    functions: Vec<FunctionDef>,

    position: usize,
    compiled: CompiledReid,
}

impl VirtualMachine {
    pub fn from(compiled: CompiledReid) -> VirtualMachine {
        VirtualMachine {
            registry_stack: Vec::new(),
            registry: Default::default(),

            stack: Vec::new(),
            heap: HashMap::new(),
            functions: Vec::new(),

            position: 0,
            compiled,
        }
    }

    pub fn add_builtin_functions<T: Into<Vec<BuiltinFunctionDef>>>(&mut self, list: T) {
        for func in list.into() {
            self.functions.push(FunctionDef::Builtin(func));
        }
    }

    pub fn run(&mut self) -> Result<(), RuntimePanic> {
        let mut error = None;
        while error.is_none() {
            if let Some(command) = self.compiled.list.get(self.position).cloned() {
                self.position += 1;
                if let Err(err) = self.execute(command.clone()) {
                    error = Some(err);
                }
                if self.remaining() == 0 {
                    break;
                }
            } else {
                error = Some(RuntimePanic::InvalidCommandIdx(self.position));
            }
        }

        if let Some(error) = error {
            Err(error)
        } else {
            Ok(())
        }
    }

    fn remaining(&self) -> usize {
        let len = self.compiled.list.len();
        len - self.position.max(0).min(len)
    }

    fn execute(&mut self, command: Command) -> Result<(), RuntimePanic> {
        match command {
            Command::InitializeVariable(heapid, vtype) => {
                self.heap.insert(heapid, AllocatedVar(vtype, None));
                Ok(())
            }
            Command::BeginScope => {
                self.registry_stack.push(self.registry.clone());
                self.registry = Default::default();
                Ok(())
            }
            Command::EndScope => {
                if let Some(reg) = self.registry_stack.pop() {
                    self.registry = reg;
                    Ok(())
                } else {
                    Err(RuntimePanic::ScopeStackUnderflow)
                }
            }
            Command::Pop(regid) => {
                self.registry[regid as usize] = Some(self.pop_stack()?);
                Ok(())
            }
            Command::Push(regid) => {
                if let Some(reg) = &self.registry[regid as usize] {
                    if self.stack.len() < usize::MAX {
                        self.stack.push(reg.clone());
                        Ok(())
                    } else {
                        Err(RuntimePanic::StackOverflow)
                    }
                } else {
                    Err(RuntimePanic::RegistryNotDefined)
                }
            }
            Command::AssignVariable(heapid, regid) => {
                if let Some(reg) = &self.registry[regid as usize] {
                    if let Some(var) = self.heap.get_mut(&heapid) {
                        var.try_set(Some(reg.clone()))?;
                        Ok(())
                    } else {
                        Err(RuntimePanic::InvalidHeapAddress)
                    }
                } else {
                    Err(RuntimePanic::RegistryNotDefined)
                }
            }
            Command::VarToReg(heapid, regid) => {
                if let Some(var) = self.heap.get(&heapid) {
                    if let Some(val) = &var.1 {
                        self.registry[regid as usize] = Some(val.clone());
                        Ok(())
                    } else {
                        Err(RuntimePanic::ValueNotInitialized)
                    }
                } else {
                    Err(RuntimePanic::InvalidHeapAddress)
                }
            }
            Command::StringLit(string) => self.push_stack(Value::StringVal(string)),
            Command::I32Lit(val) => self.push_stack(Value::I32Val(val)),
            Command::FunctionCall(funcid) => {
                if self.functions.len() <= funcid as usize {
                    Err(RuntimePanic::InvalidFuncAddress)
                } else {
                    match &self.functions[funcid as usize] {
                        FunctionDef::Builtin(f) => {
                            let mut params = Vec::new();
                            for _ in 0..f.signature.parameters.len() {
                                if let Some(val) = self.stack.pop() {
                                    params.push(val);
                                } else {
                                    return Err(RuntimePanic::StackUnderflow);
                                }
                            }
                            (f.function)(params);
                        }
                    }
                    Ok(())
                }
            }
            Command::Add => {
                let val1 = self.pop_stack()?;
                let val2 = self.pop_stack()?;
                let res = val2.add(val1);
                if let Some(res) = res {
                    self.push_stack(res);
                    Ok(())
                } else {
                    Err(RuntimePanic::AttemptedInvalidArithmeticOperation)
                }
            }
            Command::Subtract => {
                let val1 = self.pop_stack()?;
                let val2 = self.pop_stack()?;
                let res = val2.sub(val1);
                if let Some(res) = res {
                    self.push_stack(res);
                    Ok(())
                } else {
                    Err(RuntimePanic::AttemptedInvalidArithmeticOperation)
                }
            }
            Command::Mult => {
                let val1 = self.pop_stack()?;
                let val2 = self.pop_stack()?;
                let res = val2.mul(val1);
                if let Some(res) = res {
                    self.push_stack(res);
                    Ok(())
                } else {
                    Err(RuntimePanic::AttemptedInvalidArithmeticOperation)
                }
            }
        }
    }

    fn pop_stack(&mut self) -> Result<Value, RuntimePanic> {
        if let Some(val) = self.stack.pop() {
            Ok(val)
        } else {
            Err(RuntimePanic::StackUnderflow)
        }
    }

    fn push_stack(&mut self, val: Value) -> Result<(), RuntimePanic> {
        if self.stack.len() < usize::MAX {
            self.stack.push(val);
            Ok(())
        } else {
            Err(RuntimePanic::StackOverflow)
        }
    }
}

#[derive(Clone, Debug)]
struct AllocatedVar(VariableType, Option<Value>);

impl AllocatedVar {
    fn try_set(&mut self, new_val: Option<Value>) -> Result<(), RuntimePanic> {
        if let Some(val) = new_val {
            if val.get_type() == self.0 {
                self.1 = Some(val);
                Ok(())
            } else {
                Err(RuntimePanic::InvalidTypeAssign)
            }
        } else {
            self.1 = None;
            Ok(())
        }
    }
}