use std::cell::OnceCell;
use std::sync::Arc;

use crate::builtins::env;
use crate::bytecode::{BinOp, OpCode, OpCodes, Program, UnOp};
use crate::error::*;
use crate::ty::common::Symbol;
use crate::ty::internal::*;
use crate::ty::public as p;

use stack::{STACK_SIZE, Stack};

pub use env::{CapturedEnv, Env};

mod env;
mod jit;
mod stack;

#[cfg(test)]
mod test;

pub fn run(prog: Program) -> Result<p::Value> {
    let vm = VM::new(
        prog.thunks,
        prog.funcs
            .into_iter()
            .map(|f| Func {
                env: OnceCell::new(),
                param: f.param,
                opcodes: f.opcodes,
            })
            .collect(),
    );
    let env = env();
    let temp = vm.eval(prog.top_level, env)?;
    let temp = temp.to_public(&vm);
    Ok(temp)
}

pub struct VM<'vm> {
    thunks: Box<[Thunk<'vm>]>,
    funcs: Box<[Func<'vm>]>,
}

impl<'vm> VM<'vm> {
    fn new(thunks: Box<[OpCodes]>, funcs: Box<[Func<'vm>]>) -> Self {
        let thunks = thunks
            .into_iter()
            .map(|opcodes| Thunk::new(opcodes))
            .collect();
        VM { thunks, funcs }
    }

    pub fn get_thunk(&self, idx: usize) -> &'vm Thunk<'vm> {
        // SAFETY: The `idx` is within bounds as `thunks` is initialized with `prog.thunks`
        // and `idx` is expected to be a valid index into this collection.
        // The lifetime of the returned reference is tied to `&self`.
        unsafe { &*(&self.thunks[idx] as *const _) }
    }

    pub fn get_func(&self, idx: usize) -> &'vm Func<'vm> {
        // SAFETY: The `idx` is within bounds as `funcs` is initialized with `prog.funcs`
        // and `idx` is expected to be a valid index into this collection.
        // The lifetime of the returned reference is tied to `&self`.
        unsafe { &*(&self.funcs[idx] as *const _) }
    }

    pub fn eval(&self, opcodes: OpCodes, env: Arc<Env<'vm>>) -> Result<Value<'vm>> {
        let mut stack = Stack::<STACK_SIZE>::new();
        let mut iter = opcodes.into_iter();
        while let Some(opcode) = iter.next() {
            let jmp = self.single_op(opcode, &mut stack, env.clone())?;
            for _ in 0..jmp {
                iter.next().unwrap();
            }
        }
        assert_eq!(stack.len(), 1);
        let mut ret = stack.pop();
        ret.force(self)?;
        Ok(ret)
    }

    #[inline]
    fn single_op<'s, const CAP: usize>(
        &self,
        opcode: OpCode,
        stack: &'s mut Stack<'vm, CAP>,
        env: Arc<Env<'vm>>,
    ) -> Result<usize>
    {
        match opcode {
            OpCode::Illegal => panic!("illegal opcode"),
            OpCode::Const { value } => stack.push(Value::Const(value))?,
            OpCode::LoadThunk { idx } => {
                self.thunks[idx].capture(env);
                stack.push(Value::Thunk(self.thunks[idx].clone()))?
            }
            OpCode::LoadValue { idx } => {
                stack.push(self.get_thunk(idx).force(self)?)?;
            }
            OpCode::ForceValue => {
                stack.tos_mut()?.force(self)?;
            }
            OpCode::Jmp { step } => return Ok(step),
            OpCode::JmpIfTrue { step } => {
                if let Value::Const(Const::Bool(true)) = stack.pop() {
                    return Ok(step);
                }
            }
            OpCode::JmpIfFalse { step } => {
                if let Value::Const(Const::Bool(false)) = stack.pop() {
                    return Ok(step);
                }
            }
            OpCode::Call { arity } => {
                let mut args = Vec::with_capacity(arity);
                for _ in 0..arity {
                    args.insert(0, stack.pop());
                }
                let mut func = stack.pop();
                func.force(self)?;
                stack.push(func.call(self, args)?)?;
            }
            OpCode::Func { idx } => {
                let func = self.get_func(idx);
                func.env.get_or_init(|| env.captured());
                stack.push(Value::Func(func))?;
            }
            OpCode::UnOp { op } => {
                use UnOp::*;
                let mut value = stack.pop();
                value.force(self)?;
                stack.push(match op {
                    Not => value.not(),
                })?;
            }
            OpCode::BinOp { op } => {
                use BinOp::*;
                let mut rhs = stack.pop();
                let mut lhs = stack.pop();
                lhs.force(self)?;
                rhs.force(self)?;
                stack.push(match op {
                    Add => lhs.add(rhs),
                    And => lhs.and(rhs),
                    Or => lhs.or(rhs),
                    Eq => lhs.eq(rhs),
                    Con => lhs.concat(rhs),
                    Upd => lhs.update(rhs),
                })?;
            }
            OpCode::ConcatString => {
                let mut rhs = stack.pop();
                rhs.force(self)?;
                stack.tos_mut()?.concat_string(rhs);
            }
            OpCode::Path => {
                todo!()
            }
            OpCode::List => {
                stack.push(Value::List(List::empty()))?;
            }
            OpCode::PushElem => {
                let elem = stack.pop();
                stack.tos_mut()?.push(elem);
            }
            OpCode::AttrSet => {
                stack.push(Value::AttrSet(AttrSet::empty()))?;
            }
            OpCode::RecAttrSet => {
                let new = env.enter_rec();
                stack.push(Value::RecAttrSet(RecAttrSet::new(new)))?;
            }
            OpCode::PushStaticAttr { name } => {
                let val = stack.pop();
                stack.tos_mut()?.push_attr(Symbol::new(name), val);
            }
            OpCode::PushDynamicAttr => {
                let val = stack.pop();
                let mut sym = stack.pop();
                sym.force(self)?.coerce_to_string();
                let sym = sym.unwrap_const().unwrap_string().into();
                stack.tos_mut()?.push_attr(sym, val);
            }
            OpCode::Select { sym } => {
                stack.tos_mut()?.force(self)?.select(Symbol::new(sym))?;
            }
            OpCode::SelectOrDefault { sym } => {
                let default = stack.pop();
                stack
                    .tos_mut()?
                    .force(self)?
                    .select_with_default(Symbol::new(sym), default)?;
            }
            OpCode::SelectDynamic => {
                let mut val = stack.pop();
                val.force(self)?;
                val.coerce_to_string();
                let sym = val.unwrap_const().unwrap_string().into();
                stack.tos_mut()?.force(self)?.select(sym)?;
            }
            OpCode::SelectDynamicOrDefault => {
                let default = stack.pop();
                let mut val = stack.pop();
                val.force(self)?;
                val.coerce_to_string();
                let sym = val.unwrap_const().unwrap_string().into();
                stack
                    .tos_mut()?
                    .force(self)?
                    .select_with_default(sym, default)?;
            }
            OpCode::HasAttr { sym } => {
                stack.tos_mut()?.force(self)?.has_attr(Symbol::new(sym));
            }
            OpCode::HasDynamicAttr => {
                let mut val = stack.pop();
                val.coerce_to_string();
                let sym = val.unwrap_const().unwrap_string().into();
                stack.tos_mut()?.force(self)?.has_attr(sym);
            }
            OpCode::LookUp { sym } => {
                stack.push(
                    env.lookup(Symbol::new(sym.clone()))
                        .ok_or_else(|| Error::EvalError(format!(r#""{sym}" not found"#)))?,
                )?;
            }
            OpCode::EnterEnv => match stack.pop() {
                Value::AttrSet(attrs) => env.enter(attrs.into_inner()),
                Value::RecAttrSet(attrs) => env.enter(attrs.into_inner()),
                _ => unreachable!(),
            },
            OpCode::LeaveEnv => {
                env.leave();
            }
            OpCode::Assert => {
                if !stack.pop().unwrap_const().unwrap_bool() {
                    todo!()
                }
            }
        }
        Ok(0)
    }
}