use crate::instruction::Instruction;
use crate::vm::constant::Constant;
use crate::vm::function::Function;
use crate::vm::value::Value;
use std::collections::HashMap;
use std::rc::Rc;

pub mod constant;
pub mod function;
pub mod value;

pub struct DvmContext {
    functions: HashMap<Rc<str>, Function>,
    constants: HashMap<Rc<str>, Constant>,
}

impl DvmContext {
    pub fn new() -> Self {
        Self {
            functions: HashMap::new(),
            constants: HashMap::new(),
        }
    }

    pub fn add_function(&mut self, function: Function) {
        self.functions.insert(function.name_owned(), function);
    }

    pub fn constants(&self) -> &HashMap<Rc<str>, Constant> {
        &self.constants
    }

    pub fn add_constant(&mut self, constant: Constant) {
        match &constant {
            Constant::String(string_constant) => {
                self.constants
                    .insert(string_constant.name_owned(), constant);
            }
        }
    }
}

pub struct DvmState {
    stack: Vec<Value>,
    registers: Vec<Value>,
    ip: usize,
    fp: usize,
}

impl DvmState {
    pub fn new() -> Self {
        Self {
            stack: vec![],
            registers: vec![],
            ip: 0,
            fp: 0,
        }
    }

    pub fn stack(&self) -> &Vec<Value> {
        &self.stack
    }

    pub fn stack_mut(&mut self) -> &mut Vec<Value> {
        &mut self.stack
    }

    pub fn registers(&self) -> &Vec<Value> {
        &self.registers
    }

    pub fn registers_mut(&mut self) -> &mut Vec<Value> {
        &mut self.registers
    }

    pub fn ensure_registers(&mut self, count: usize) {
        self.registers.resize_with(count, Default::default);
    }

    pub fn ip(&self) -> usize {
        self.ip
    }

    pub fn increment_ip(&mut self) {
        self.ip += 1;
    }

    pub fn fp(&self) -> usize {
        self.fp
    }

    pub fn set_fp(&mut self, fp: usize) {
        self.fp = fp;
    }
}

pub fn call(
    context: &DvmContext,
    state: &mut DvmState,
    function_name: &str,
    arguments: Vec<Value>,
) -> Option<Value> {
    let function = context
        .functions
        .get(function_name)
        .expect(&format!("Function {} not found", function_name));

    let instructions = function.instructions();
    state.ensure_registers(function.register_count());

    // put each arg on the stack
    for argument in arguments {
        state.stack_mut().push(argument);
    }

    while state.ip() < instructions.len() {
        let instruction = &instructions[state.ip()];
        match instruction {
            /* Move instructions */
            Instruction::MoveRegister(source, destination) => {
                // copy value from one register to another register
                let value = state.registers()[*source].clone();
                state.registers_mut()[*destination] = value;
            }
            Instruction::MoveInt(value, destination) => {
                state.registers_mut()[*destination] = Value::Int(*value);
            }
            Instruction::MoveStackFrameOffset(offset, destination) => {
                // copy a value offset from the current frame pointer (fp) to a register
                let value_index = state
                    .fp()
                    .checked_add_signed(*offset)
                    .expect("Overflow when adding offset to fp");
                let value = state.stack()[value_index].clone();
                state.registers_mut()[*destination] = value;
            }

            /* Push instructions */
            Instruction::PushRegister(source) => {
                // copy a value from a register to the top of the stack
                let value = state.registers()[*source].clone();
                state.stack_mut().push(value);
            }
            Instruction::PushInt(value) => {
                state.stack_mut().push(Value::Int(*value));
            }
            Instruction::PushStackFrameOffset(offset) => {
                // copy a value from somewhere on the stack to the top of the stack
                let value_index = state
                    .fp()
                    .checked_add_signed(*offset)
                    .expect("Overflow when adding offset to fp");
                let value = state.stack()[value_index].clone();
                state.stack_mut().push(value);
            }

            /* Invoke instructions */
            Instruction::InvokePlatformStatic(function_name) => {
                if function_name.as_ref() == "println" {
                    println!("{:?}", state.stack());
                    println!("{:?}", state.registers());
                }
            }

            /* Load constant instructions */
            Instruction::LoadStringConstant(constant_name, destination) => {
                let constant = &context.constants()[constant_name];
                match constant {
                    Constant::String(string_constant) => {
                        state.registers_mut()[*destination] =
                            Value::String(string_constant.content_owned());
                    }
                }
            }

            /* Pop instructions */
            Instruction::Pop(maybe_register) => {
                let value = state.stack_mut().pop().unwrap();
                if let Some(register) = maybe_register {
                    state.registers_mut()[*register] = value;
                }
            }
        }
        state.increment_ip();
    }
    None
}