mod array;
pub mod constant;
mod field;
pub mod function;
pub mod implementation;
pub mod instruction;
mod interface;
pub mod lib;
mod method;
pub mod object;
pub mod op_codes;
pub mod platform;
pub mod source_code_location;
pub mod r#type;
pub mod util;
pub mod value;
mod virtual_method;

use crate::vm::array::DvmArray;
use crate::vm::constant::{DvmConstant, DvmConstantArray};
use crate::vm::implementation::DvmImplementation;
use crate::vm::instruction::{Immediate, Instruction, Location};
use crate::vm::object::DvmObject;
use crate::vm::value::DvmValue;
use function::DvmFunction;
use source_code_location::SourceCodeLocation;
use std::cell::RefCell;
use std::collections::HashMap;
use std::fmt::Display;
use std::rc::Rc;

pub fn dump_state(state: &DvmState, message: &str) {
    println!("----");
    println!("{}", message);
    println!("----");
    println!(
        "Current Instruction: {:?}",
        state.current_instruction.clone().unwrap()
    );
    println!("Registers:");
    for (i, register) in state.registers.iter().enumerate() {
        println!("\tr{:x}: {:?}", i, register);
    }
    println!("Call stack:");
    for call in state.call_stack.iter() {
        println!("\t{}", call);
    }
    println!("Stack:");
    for (i, value) in state.stack.iter().enumerate() {
        println!("\t{}: {}", i, value);
    }
    println!("----");
}

macro_rules! dvm_panic {
    ($context: expr, $message: expr) => {
        dump_state($context, $message);
        panic!()
    };
}

pub type PlatformFunction = fn(state: &mut DvmState, context: &DvmContext);

pub struct DvmState {
    call_stack: Vec<DvmCall>,
    current_instruction: Option<Instruction>,
    registers: Vec<DvmValue>,
    stack: Vec<DvmValue>,
}

impl DvmState {
    pub fn new() -> Self {
        DvmState {
            call_stack: vec![],
            current_instruction: None,
            registers: vec![DvmValue::Empty; 16],
            stack: vec![],
        }
    }

    pub fn pop_stack(&mut self) -> DvmValue {
        self.stack.pop().unwrap_or_else(|| {
            dvm_panic!(self, "Stack underflow!");
        })
    }

    pub fn get_stack_argument(&self, index: usize) -> DvmValue {
        self.stack
            .get(self.stack.len() - 2 - index)
            .unwrap_or_else(|| {
                dvm_panic!(self, "Stack underflow!");
            })
            .clone()
    }
}

pub struct DvmContext {
    main_function: Rc<DvmFunction>,
    implementations: HashMap<String, Rc<DvmImplementation>>,
    static_functions: HashMap<String, Rc<DvmFunction>>,
    platform_functions: HashMap<String, PlatformFunction>,
}

impl DvmContext {
    pub fn new(main_function: Rc<DvmFunction>) -> Self {
        DvmContext {
            main_function,
            implementations: HashMap::new(),
            static_functions: HashMap::new(),
            platform_functions: HashMap::new(),
        }
    }

    pub fn add_platform_functions(
        &mut self,
        platform_functions: HashMap<String, PlatformFunction>,
    ) {
        for (fqn, platform_function) in platform_functions {
            self.platform_functions.insert(fqn, platform_function);
        }
    }
}

pub struct DvmCall {
    function_name: Rc<String>,
    source_code_location: SourceCodeLocation,
}

impl Display for DvmCall {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_fmt(format_args!(
            "{} in {} @ {}:{}",
            self.function_name,
            self.source_code_location.source_file_name,
            self.source_code_location.line,
            self.source_code_location.char
        ))
    }
}

fn update_index(index: &mut usize, offset: isize) {
    if offset.is_negative() {
        *index -= offset.abs() as usize;
    } else {
        *index += offset as usize;
    }
}

fn push_call_frame(
    function_name: Rc<String>,
    source_code_location: SourceCodeLocation,
    state: &mut DvmState,
) {
    state.call_stack.push(DvmCall {
        function_name,
        source_code_location,
    });
}

fn pop_call_frame(state: &mut DvmState) {
    state.call_stack.pop();
}

fn compute_index_with_offset(
    registers: &Vec<DvmValue>,
    index_register: usize,
    offset: &Option<isize>,
) -> usize {
    let index = registers[index_register].as_usize().expect(&format!(
        "Could not convert the value of index_register {} to usize; found {:?}",
        index_register, registers[index_register]
    ));
    if let Some(o) = offset {
        if o.is_negative() {
            index - o.unsigned_abs()
        } else {
            index + o.unsigned_abs()
        }
    } else {
        index
    }
}

fn copy_from_source(state: &DvmState, source: &Location) -> DvmValue {
    match source {
        Location::Register(register) => state.registers[*register].clone(),
        Location::Stack { offset } => state.stack[*offset].clone(),
        Location::Field {
            object_register,
            field_name,
        } => state.registers[*object_register].map_object(
            |o| o.borrow().read_field(field_name).clone(),
            |v| {
                dvm_panic!(state, &format!("Expected an object, but found: {}", v));
            },
        ),
        Location::Array {
            array_register,
            index_register,
            offset,
        } => {
            let index = compute_index_with_offset(&state.registers, *index_register, offset);
            state.registers[*array_register].map_array(
                |a| a.borrow().read_element(index),
                |v| {
                    dvm_panic!(state, &format!("Expected an array, but found: {}", v));
                },
            )
        }
    }
}

fn write_to_destination(state: &mut DvmState, destination: &Location, value: DvmValue) {
    match destination {
        Location::Register(register) => {
            state.registers[*register] = value;
        }
        Location::Stack { offset } => {
            state.stack[*offset] = value;
        }
        Location::Field {
            object_register,
            field_name,
        } => {
            state.registers[*object_register]
                .expect_object()
                .borrow_mut()
                .write_field(field_name, value);
        }
        Location::Array {
            array_register,
            index_register,
            offset,
        } => {
            let index = compute_index_with_offset(&state.registers, *index_register, offset);
            state.registers[*array_register]
                .expect_array()
                .borrow_mut()
                .write_element(index, value);
        }
    }
}

fn constant_array_to_array(constant_array: &DvmConstantArray) -> DvmArray {
    match constant_array {
        DvmConstantArray::Bytes(v) => DvmArray::Bytes(v.clone()),
        DvmConstantArray::Ints(v) => DvmArray::Ints(v.clone()),
        DvmConstantArray::Longs(v) => DvmArray::Longs(v.clone()),
        DvmConstantArray::Doubles(v) => DvmArray::Doubles(v.clone()),
        DvmConstantArray::USizes(v) => DvmArray::USizes(v.clone()),
        DvmConstantArray::Booleans(v) => DvmArray::Booleans(v.clone()),
        DvmConstantArray::Strings(vs) => {
            DvmArray::Strings(vs.iter().map(|s| Rc::new(s.clone())).collect())
        }
        DvmConstantArray::Arrays(va) => DvmArray::Arrays(
            va.iter()
                .map(|a| Rc::new(RefCell::new(constant_array_to_array(a))))
                .collect(),
        ),
    }
}

fn constant_to_value(constant: &DvmConstant) -> DvmValue {
    match constant {
        DvmConstant::String(s) => DvmValue::String(Rc::new(s.clone())),
        DvmConstant::Array(a) => DvmValue::Array(Rc::new(RefCell::new(constant_array_to_array(a)))),
    }
}

fn immediate_to_value(immediate: &Immediate) -> DvmValue {
    match immediate {
        Immediate::Byte(b) => DvmValue::Byte(*b),
        Immediate::Int(i) => DvmValue::Int(*i),
        Immediate::Long(l) => DvmValue::Long(*l),
        Immediate::Double(d) => DvmValue::Double(*d),
        Immediate::Usize(s) => DvmValue::USize(*s),
        Immediate::Boolean(b) => DvmValue::Boolean(*b),
        Immediate::Constant(c) => constant_to_value(c),
        Immediate::Empty => DvmValue::Empty,
    }
}

pub fn run_main_function(state: &mut DvmState, context: &DvmContext) -> i32 {
    let main_function = context.main_function.clone();
    push_call_frame(
        Rc::new(main_function.fqn().to_string()),
        main_function.source_code_location().clone(),
        state,
    );
    run(main_function.instructions(), state, context);
    pop_call_frame(state);
    state.pop_stack().expect_int_or_else(|v| {
        dvm_panic!(state, &format!("Expected DvmValue::Int, but found {}", v));
    })
}

// TODO: find all possible unwraps/expects and wrap with call to dvm_panic
pub fn run(instructions: &[Instruction], state: &mut DvmState, context: &DvmContext) {
    let mut index = 0;
    while index < instructions.len() {
        state.current_instruction = Some(instructions[index].clone());
        use Instruction::*;
        match &instructions[index] {
            MoveImmediate {
                destination,
                immediate,
            } => {
                write_to_destination(state, destination, immediate_to_value(immediate));
            }
            Copy {
                source,
                destination,
            } => {
                let value = copy_from_source(state, source);
                write_to_destination(state, destination, value);
            }
            Push { source_register } => {
                state.stack.push(state.registers[*source_register].clone());
            }
            Pop {
                destination_register,
            } => {
                let value = state.stack.pop().unwrap_or_else(|| {
                    dvm_panic!(state, "Stack underflow");
                });
                if let Some(register_index) = destination_register {
                    state.registers[*register_index] = value;
                }
            }
            AllocateObject {
                implementation_name,
                destination_register,
            } => {
                let implementation = context
                    .implementations
                    .get(implementation_name.as_str())
                    .unwrap_or_else(|| {
                        dvm_panic!(
                            state,
                            &format!("Cannot find implementation {}", implementation_name)
                        );
                    })
                    .clone();
                let object = DvmObject::new(implementation);
                state.registers[*destination_register] =
                    DvmValue::Object(Rc::new(RefCell::new(object)));
            }
            InvokeStatic {
                function_name,
                source_code_location,
            } => {
                // Find function
                let static_function = context
                    .static_functions
                    .get(function_name.as_str())
                    .unwrap_or_else(|| {
                        dvm_panic!(
                            state,
                            &format!("Cannot find static function {}", function_name)
                        );
                    })
                    .clone();

                // Do call
                state.stack.push(DvmValue::Empty); // space for return value
                push_call_frame(function_name.clone(), source_code_location.clone(), state);
                run(static_function.instructions(), state, context);
                pop_call_frame(state);
            }
            InvokeObject {
                object_register,
                function_name,
                source_code_location,
            } => {
                // Find method and get instructions
                let object_value = state.registers[*object_register].clone();
                let object_rc = object_value.expect_object_or_else(|v| {
                    dvm_panic!(
                        state,
                        &format!("Expected DvmValue::Object, but found {}", object_value)
                    );
                });
                let object_ref = object_rc.borrow();
                let method = object_ref
                    .get_method(function_name.as_str())
                    .unwrap_or_else(|| {
                        dvm_panic!(
                            state,
                            &format!(
                                "Cannot find method {} for object {:?}",
                                function_name, object_ref
                            )
                        );
                    });
                let function = method.dm_fn();
                let instructions = function.instructions();

                // Do call
                state.stack.push(DvmValue::Empty); // space for return value
                state.stack.push(object_value); // self object
                push_call_frame(function_name.clone(), source_code_location.clone(), state);
                run(instructions, state, context);
                pop_call_frame(state);
                state.stack.pop(); // self object
            }
            InvokeStaticPlatform {
                function_name,
                source_code_location,
            } => {
                // Find function
                let platform_function = context
                    .platform_functions
                    .get(function_name.as_str())
                    .unwrap_or_else(|| {
                        dvm_panic!(
                            state,
                            &format!("Cannot find static platform function {}", function_name)
                        );
                    });

                // Do call
                state.stack.push(DvmValue::Empty); // space for return value
                push_call_frame(function_name.clone(), source_code_location.clone(), state);
                platform_function(state, context);
                pop_call_frame(state);
            }
            InvokeObjectPlatform {
                object_register,
                function_name,
                source_code_location,
            } => {
                // Find function
                let object_platform_function = context
                    .platform_functions
                    .get(function_name.as_str())
                    .unwrap_or_else(|| {
                        dvm_panic!(
                            state,
                            &format!("Cannot find object platform function {}", function_name)
                        );
                    });

                // Get self object
                let object_value = state.registers[*object_register].clone();
                if !object_value.is_object() {
                    dvm_panic!(
                        state,
                        &format!(
                            "Expected DvmValue::Object, but found DvmValue::{:?}",
                            object_value
                        )
                    );
                }

                // Do call
                state.stack.push(DvmValue::Empty); // space for return value
                state.stack.push(object_value); // self object
                push_call_frame(function_name.clone(), source_code_location.clone(), state);
                object_platform_function(state, context);
                pop_call_frame(state);
                state.stack.pop(); // self object
            }
            Add { .. } => {}
            Subtract { .. } => {}
            Multiply { .. } => {}
            Divide { .. } => {}
            Modulo { .. } => {}
            Power { .. } => {}
            Jump { offset } => {
                update_index(&mut index, *offset);
            }
            JumpEqual {
                left_register,
                right_register,
                offset,
            } => {
                let left_value = &state.registers[*left_register];
                let right_value = &state.registers[*right_register];
                if left_value == right_value {
                    update_index(&mut index, *offset);
                }
            }
            JumpNotEqual { .. } => {}
            JumpLessThan { .. } => {}
            JumpGreaterThan { .. } => {}
            JumpLessThanEqual { .. } => {}
            JumpGreaterThanEqual { .. } => {}
            Return => {
                break;
            }
        }
        index += 1;
    }
}