use std::cell::{Cell, RefCell};
use std::rc::Rc;

use crate::compiler::{
    Export, Function, Instruction, Module, EXTERN_BUILTIN_LIST_GET_ITERATOR,
    EXTERN_BUILTIN_LIST_ITERATOR_NEXT, EXTERN_BUILTIN_NOOP,
};
use crate::semantics::Type;
use thiserror::Error;

#[derive(Error, Debug)]
pub enum VMErrorCode {
    #[error("code panic (syntax or semantic error): {0}")]
    Panic(Rc<str>),
    #[error("internal error: stack underflow")]
    StackUnderflow,
    #[error("internal error: stack type mismatch: {0:?} is not {1:?}")]
    StackTypeMismatch(StackValue, Type),

    // TODO: This one is *not* like the others! Distinguish between internal
    //       errors and user errors?
    #[error("divide by zero")]
    DivideByZero,

    #[error("internal error: argument {0} out of range")]
    ArgumentOutOfRange(usize),
    #[error("internal error: global {0} out of range")]
    GlobalOutOfRange(usize),
    #[error("internal error: local {0} out of range")]
    LocalOutOfRange(usize),
    #[error("internal error: string {0} out of range")]
    StringOutOfRange(usize),
    #[error("internal error: function {0} out of range")]
    FunctionOutOfRange(usize),
    #[error("internal error: stack type mismatch ({0:?} is not function)")]
    StackExpectedFunction(StackValue),
    #[error("internal error: stack type mismatch ({0:?} is not object)")]
    StackExpectedObject(StackValue),
    #[error("internal error: stack type mismatch ({0:?} is not list)")]
    StackExpectedList(StackValue),
    #[error("internal error: stack type mismatch ({0:?} is not list iterator)")]
    StackExpectedListIterator(StackValue),
    #[error("internal error: slot {0} was out of range for object (type {1} with {2} slots)")]
    SlotOutOfRange(usize, Rc<str>, usize),
    #[error("internal error: the extern function with ID {0} was not registered")]
    UnregisteredExternFunction(usize),
}

#[derive(Debug)]
pub struct VMError {
    pub code: VMErrorCode,
    pub stack: Box<[Frame]>,
}

type Result<T> = std::result::Result<T, VMErrorCode>;

#[derive(Clone, Debug)]
pub struct Object {
    name: Rc<str>,
    class_id: i64,
    values: RefCell<Box<[StackValue]>>,
}

impl Object {
    pub fn get_slot(&self, index: usize) -> Result<StackValue> {
        match self.values.borrow().get(index) {
            Some(v) => Ok(v.clone()),
            None => Err(VMErrorCode::SlotOutOfRange(
                index,
                self.name.clone(),
                self.values.borrow().len(),
            )),
        }
    }
}

#[derive(Clone, Debug)]
pub struct ListIterator {
    list: Rc<Vec<StackValue>>,
    index: Cell<usize>,
}

#[derive(Clone, Debug)]
pub enum StackValue {
    Nothing,
    Bool(bool),
    Float(f64),
    Int(i64),
    String(Rc<str>),
    Function(Rc<Function>),
    ExternFunction(usize),
    Object(Rc<Object>),
    List(Rc<Vec<StackValue>>),
    ListIterator(Rc<ListIterator>),
}

impl StackValue {
    pub fn is_object(&self, id: i64) -> bool {
        match self {
            StackValue::Object(o) => o.class_id == id,
            _ => false,
        }
    }

    pub fn is_float(&self) -> bool {
        match self {
            StackValue::Float(_) => true,
            _ => false,
        }
    }

    pub fn is_nothing(&self) -> bool {
        match self {
            StackValue::Nothing => true,
            _ => false,
        }
    }

    pub fn is_bool(&self) -> bool {
        match self {
            StackValue::Bool(_) => true,
            _ => false,
        }
    }

    pub fn is_string(&self) -> bool {
        match self {
            StackValue::String(_) => true,
            _ => false,
        }
    }
}

enum FuncValue {
    Function(Rc<Function>),
    ExternFunction(usize),
}

#[derive(Debug)]
pub struct Frame {
    func: Rc<Function>,
    args: Vec<StackValue>,
    locals: Vec<StackValue>,
    stack: Vec<StackValue>,
    pc: usize,
}

impl Frame {
    fn from_function(func: Rc<Function>, args: Vec<StackValue>) -> Self {
        let mut locals = Vec::new();
        locals.resize(func.locals(), StackValue::Nothing);

        Frame {
            func,
            args,
            locals,
            stack: Vec::new(),
            pc: 0,
        }
    }

    pub fn func(&self) -> Rc<Function> {
        self.func.clone()
    }

    pub fn args(&self) -> &[StackValue] {
        &self.args
    }

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

    fn pop_value(&mut self) -> Result<StackValue> {
        self.stack.pop().ok_or_else(|| VMErrorCode::StackUnderflow)
    }

    fn pop_bool(&mut self) -> Result<bool> {
        match self.pop_value()? {
            StackValue::Bool(v) => Ok(v),
            v => Err(VMErrorCode::StackTypeMismatch(v, Type::Bool)),
        }
    }

    fn pop_float(&mut self) -> Result<f64> {
        match self.pop_value()? {
            StackValue::Float(v) => Ok(v),
            v => Err(VMErrorCode::StackTypeMismatch(v, Type::F64)),
        }
    }

    fn pop_string(&mut self) -> Result<Rc<str>> {
        match self.pop_value()? {
            StackValue::String(v) => Ok(v),
            v => Err(VMErrorCode::StackTypeMismatch(v, Type::String)),
        }
    }

    fn pop_object(&mut self) -> Result<Rc<Object>> {
        match self.pop_value()? {
            StackValue::Object(v) => Ok(v),
            v => Err(VMErrorCode::StackExpectedObject(v)),
        }
    }

    fn pop_int(&mut self) -> Result<i64> {
        match self.pop_value()? {
            StackValue::Int(v) => Ok(v),
            v => Err(VMErrorCode::StackTypeMismatch(v, Type::I64)),
        }
    }

    fn push_value(&mut self, v: StackValue) {
        self.stack.push(v)
    }

    fn push_bool(&mut self, value: bool) {
        self.push_value(StackValue::Bool(value));
    }

    fn push_float(&mut self, value: f64) {
        self.push_value(StackValue::Float(value));
    }

    fn push_nothing(&mut self) {
        self.push_value(StackValue::Nothing);
    }

    fn push_string(&mut self, v: Rc<str>) {
        self.push_value(StackValue::String(v))
    }

    fn push_function(&mut self, v: Rc<Function>) {
        self.push_value(StackValue::Function(v));
    }

    fn push_extern_function(&mut self, v: usize) {
        self.push_value(StackValue::ExternFunction(v));
    }

    fn push_object(&mut self, v: Rc<Object>) {
        self.push_value(StackValue::Object(v));
    }

    fn push_int(&mut self, v: i64) {
        self.push_value(StackValue::Int(v));
    }

    fn push_list(&mut self, value: Rc<Vec<StackValue>>) {
        self.push_value(StackValue::List(value));
    }

    fn get_argument(&self, i: usize) -> Result<StackValue> {
        self.args
            .get(i)
            .map(|v| v.clone())
            .ok_or_else(|| VMErrorCode::ArgumentOutOfRange(i))
    }

    fn get_local(&self, i: usize) -> Result<StackValue> {
        self.locals
            .get(i)
            .map(|v| v.clone())
            .ok_or_else(|| VMErrorCode::LocalOutOfRange(i))
    }

    fn get_string(&self, i: usize) -> Result<Rc<str>> {
        let strings = self.func.strings();
        strings
            .get(i)
            .map(|v| v.clone())
            .ok_or_else(|| VMErrorCode::StringOutOfRange(i))
    }

    fn store_local(&mut self, i: usize, v: StackValue) -> Result<()> {
        if i >= self.locals.len() {
            Err(VMErrorCode::LocalOutOfRange(i))
        } else {
            self.locals[i] = v;
            Ok(())
        }
    }

    fn store_argument(&mut self, i: usize, v: StackValue) -> Result<()> {
        if i >= self.locals.len() {
            Err(VMErrorCode::LocalOutOfRange(i))
        } else {
            self.args[i] = v;
            Ok(())
        }
    }

    fn pop_function(&mut self) -> Result<FuncValue> {
        match self.pop_value()? {
            StackValue::Function(f) => Ok(FuncValue::Function(f)),
            StackValue::ExternFunction(i) => Ok(FuncValue::ExternFunction(i)),
            v => Err(VMErrorCode::StackExpectedFunction(v)),
        }
    }
}

pub struct Context {
    module: Rc<Module>,
    globals: Vec<StackValue>,
}

impl Context {
    pub fn new(module: Rc<Module>) -> Context {
        let mut globals = Vec::new();
        globals.resize(module.globals, StackValue::Nothing);
        Context { module, globals }
    }

    pub fn init(&mut self) -> std::result::Result<(), VMError> {
        let init_index = self.module.init;
        let init_fn = self.module.functions[init_index].clone();
        eval(self, init_fn, vec![])?;
        Ok(())
    }

    fn get_global(&self, i: usize) -> Result<StackValue> {
        self.globals
            .get(i)
            .map(|v| v.clone())
            .ok_or_else(|| VMErrorCode::GlobalOutOfRange(i)) // TODO: Test
    }

    fn set_global(&mut self, i: usize, v: StackValue) -> Result<()> {
        if i >= self.globals.len() {
            Err(VMErrorCode::GlobalOutOfRange(i)) // TODO: Test
        } else {
            self.globals[i] = v;
            Ok(())
        }
    }

    fn get_function(&self, i: usize) -> Result<Rc<Function>> {
        let functions = self.module.functions();
        functions
            .get(i)
            .map(|v| v.clone())
            .ok_or_else(|| VMErrorCode::FunctionOutOfRange(i)) // TODO: Test
    }

    fn call_extern_function(&self, index: usize, args: &[StackValue]) -> Result<StackValue> {
        match index {
            EXTERN_BUILTIN_NOOP => Ok(StackValue::Nothing),
            EXTERN_BUILTIN_LIST_GET_ITERATOR => {
                let Some(list_value) = args.get(0) else {
                    return Err(VMErrorCode::ArgumentOutOfRange(0));
                };
                let StackValue::List(list) = list_value else {
                    return Err(VMErrorCode::StackExpectedList(list_value.clone()));
                };

                Ok(StackValue::ListIterator(Rc::new(ListIterator {
                    list: list.clone(),
                    index: Cell::new(0),
                })))
            }
            EXTERN_BUILTIN_LIST_ITERATOR_NEXT => {
                let Some(iter_value) = args.get(0) else {
                    return Err(VMErrorCode::ArgumentOutOfRange(0));
                };
                let StackValue::ListIterator(li) = iter_value else {
                    return Err(VMErrorCode::StackExpectedListIterator(iter_value.clone()));
                };

                let index = li.index.get();
                if index >= li.list.len() {
                    Ok(StackValue::Nothing)
                } else {
                    let result = li.list[index].clone();
                    li.index.set(index + 1);
                    Ok(result)
                }
            }
            _ => Err(VMErrorCode::UnregisteredExternFunction(index)),
        }
    }
}

enum Flow {
    Break,
    Continue,
}

#[inline(always)]
fn eval_one(
    instruction: Instruction,
    index: &mut usize,
    c: &mut Context,
    f: &mut Frame,
    stack: &mut Vec<Frame>,
) -> Result<Flow> {
    match instruction {
        Instruction::Panic(index) => {
            let v = f
                .get_string(index)
                .unwrap_or_else(|_| "!!panic string out of range!!".into());
            return Err(VMErrorCode::Panic(v));
        }
        Instruction::BoolNot => {
            let value = f.pop_bool()?;
            f.push_bool(!value);
        }
        Instruction::Discard => {
            f.pop_value()?;
        }
        Instruction::Dup => {
            let v = f.pop_value()?;
            f.push_value(v.clone());
            f.push_value(v);
        }
        Instruction::FloatAdd => {
            let x = f.pop_float()?;
            let y = f.pop_float()?;
            f.push_float(x + y);
        }
        Instruction::FloatDivide => {
            let x = f.pop_float()?;
            let y = f.pop_float()?;
            if x == 0. {
                return Err(VMErrorCode::DivideByZero);
            }
            f.push_float(y / x);
        }
        Instruction::FloatMultiply => {
            let x = f.pop_float()?;
            let y = f.pop_float()?;
            f.push_float(x * y);
        }
        Instruction::FloatSubtract => {
            let x = f.pop_float()?;
            let y = f.pop_float()?;
            f.push_float(y - x);
        }
        Instruction::Jump(i) => {
            *index = i;
        }
        Instruction::JumpFalse(i) => {
            if !(f.pop_bool()?) {
                *index = i;
            }
        }
        Instruction::JumpTrue(i) => {
            if f.pop_bool()? {
                *index = i;
            }
        }
        Instruction::LoadArgument(i) => {
            let v = f.get_argument(i)?;
            f.push_value(v);
        }
        Instruction::LoadLocal(i) => {
            let v = f.get_local(i)?;
            f.push_value(v);
        }
        Instruction::LoadModule(i) => {
            let v = c.get_global(i)?;
            f.push_value(v);
        }
        Instruction::PushFalse => {
            f.push_bool(false);
        }
        Instruction::PushFloat(v) => {
            f.push_float(v);
        }
        Instruction::PushNothing => {
            f.push_nothing();
        }
        Instruction::PushString(s) => {
            let v = f.get_string(s)?;
            f.push_string(v);
        }
        Instruction::PushTrue => {
            f.push_bool(true);
        }
        Instruction::StoreArgument(i) => {
            let v = f.pop_value()?;
            f.store_argument(i, v)?;
        }
        Instruction::StoreLocal(i) => {
            let v = f.pop_value()?;
            f.store_local(i, v)?;
        }
        Instruction::StoreModule(i) => {
            let v = f.pop_value()?;
            c.set_global(i, v)?;
        }
        Instruction::StoreSlot(i) => {
            let o = f.pop_object()?;
            let v = f.pop_value()?;
            o.values.borrow_mut()[i] = v;
        }
        Instruction::LoadFunction(i) => {
            let v = c.get_function(i)?;
            f.push_function(v);
        }
        Instruction::LoadExternFunction(i) => {
            f.push_extern_function(i);
        }
        Instruction::Call(arg_count) => {
            let function = f.pop_function()?;
            let mut args = Vec::new();
            for _ in 0..arg_count {
                args.push(f.pop_value()?);
            }
            match function {
                FuncValue::Function(func) => {
                    let mut frame = Frame::from_function(func, args);
                    std::mem::swap(&mut frame, f);
                    frame.pc = *index;
                    stack.push(frame);
                    *index = 0;
                }
                FuncValue::ExternFunction(i) => {
                    f.push_value(c.call_extern_function(i, &args)?);
                }
            }
        }
        Instruction::Return => match stack.pop() {
            Some(mut frame) => {
                // The return value is at the top of the stack already.
                let retval = f.pop_value()?;
                std::mem::swap(&mut frame, f);
                *index = f.pc;
                f.push_value(retval);
            }
            None => return Ok(Flow::Break),
        },
        Instruction::StringAdd => {
            let x = f.pop_string()?;
            let y = f.pop_string()?;

            let mut new_string = x.to_string();
            new_string.push_str(&y);

            f.push_string(new_string.into());
        }
        Instruction::EqBool => {
            let x = f.pop_bool()?;
            let y = f.pop_bool()?;
            f.push_bool(x == y);
        }
        Instruction::EqFloat => {
            let x = f.pop_float()?;
            let y = f.pop_float()?;
            f.push_bool(x == y);
        }
        Instruction::EqString => {
            let x = f.pop_string()?;
            let y = f.pop_string()?;
            f.push_bool(x == y);
        }
        Instruction::GreaterFloat => {
            let x = f.pop_float()?;
            let y = f.pop_float()?;
            f.push_bool(y > x);
        }
        Instruction::GreaterString => {
            let x = f.pop_string()?;
            let y = f.pop_string()?;
            f.push_bool(y > x);
        }
        Instruction::LessFloat => {
            let x = f.pop_float()?;
            let y = f.pop_float()?;
            f.push_bool(y < x);
        }
        Instruction::LessString => {
            let x = f.pop_string()?;
            let y = f.pop_string()?;
            f.push_bool(y < x);
        }

        Instruction::NewObject(slots) => {
            let class_id = f.pop_int()?;
            let name = f.pop_string()?;
            let mut values = Vec::with_capacity(slots);
            for _ in 0..slots {
                values.push(f.pop_value()?);
            }

            let object = Object {
                name,
                class_id,
                values: RefCell::new(values.into()),
            };

            f.push_object(object.into());
        }
        Instruction::LoadSlot(slot) => {
            let obj = f.pop_object()?;
            f.push_value(obj.get_slot(slot)?);
        }
        Instruction::IsClass(id) => {
            let value = f.pop_value()?;
            f.push_bool(value.is_object(id));
        }
        Instruction::PushInt(v) => {
            f.push_int(v);
        }
        Instruction::IsBool => {
            let v = f.pop_value()?;
            f.push_bool(v.is_bool());
        }
        Instruction::IsFloat => {
            let v = f.pop_value()?;
            f.push_bool(v.is_float());
        }
        Instruction::IsString => {
            let v = f.pop_value()?;
            f.push_bool(v.is_string());
        }
        Instruction::IsNothing => {
            let v = f.pop_value()?;
            f.push_bool(v.is_nothing());
        }
        Instruction::NewList(c) => {
            let mut v = Vec::with_capacity(c);
            for _ in 0..c {
                v.push(f.pop_value()?);
            }
            f.push_list(Rc::new(v));
        }
    }

    Ok(Flow::Continue)
}

pub fn eval(
    c: &mut Context,
    function: Rc<Function>,
    args: Vec<StackValue>,
) -> std::result::Result<StackValue, VMError> {
    let mut stack = Vec::new();
    let mut f = Frame::from_function(function, args);

    let mut index = 0;
    loop {
        let instructions = f.func.instructions();
        let instruction = instructions[index];

        // {
        //     eprint!("{index}: [");
        //     if f.stack.len() > 3 {
        //         eprint!("...");
        //     }
        //     for val in f.stack.iter().take(3) {
        //         eprint!("{val:?} ");
        //     }
        //     eprintln!("] => {instruction:?}");
        // }

        index += 1;

        match eval_one(instruction, &mut index, c, &mut f, &mut stack) {
            Ok(Flow::Break) => match f.pop_value() {
                Ok(v) => return Ok(v),
                Err(e) => {
                    f.pc = index;
                    stack.push(f);
                    return Err(VMError {
                        code: e,
                        stack: stack.into(),
                    });
                }
            },
            Ok(Flow::Continue) => (),
            Err(e) => {
                f.pc = index;
                stack.push(f);
                return Err(VMError {
                    code: e,
                    stack: stack.into(),
                });
            }
        };
    }
}

pub fn eval_export_fn(
    c: &mut Context,
    name: &str,
    args: &[StackValue],
) -> std::result::Result<StackValue, VMError> {
    let export = match c.module.exports.get(name) {
        Some(Export::Function(id)) => id,
        Some(_) => todo!(),
        None => todo!(),
    };

    let function = c.module.functions[*export].clone();
    let args = args.iter().map(|a| a.clone()).collect();
    eval(c, function, args)
}