use std::cell::{Cell, OnceCell};

use derive_more::Constructor;
use hashbrown::HashMap;
use inkwell::execution_engine::JitFunction;
use itertools::Itertools;

use crate::bytecode::Func as BFunc;
use crate::error::Result;
use crate::ir;
use crate::jit::JITFunc;
use crate::ty::internal::{Thunk, Value};
use crate::vm::{VM, VmEnv};

#[derive(Debug, Clone)]
pub enum Param {
    Ident(usize),
    Formals {
        formals: Vec<(usize, Option<usize>)>,
        ellipsis: bool,
        alias: Option<usize>,
    },
}

impl From<ir::Param> for Param {
    fn from(value: ir::Param) -> Self {
        match value {
            ir::Param::Ident(ident) => Param::Ident(ident),
            ir::Param::Formals {
                formals,
                ellipsis,
                alias,
            } => Param::Formals {
                formals: formals
                    .into_iter()
                    .map(|(sym, default)| (sym, default.map(|default| default.idx)))
                    .collect(),
                ellipsis,
                alias,
            },
        }
    }
}

#[derive(Debug, Clone, Constructor)]
pub struct Func<'jit: 'vm, 'vm> {
    pub func: &'vm BFunc,
    pub env: &'vm VmEnv<'jit, 'vm>,
    pub compiled: OnceCell<JitFunction<'jit, JITFunc<'jit, 'vm>>>,
    pub count: Cell<usize>,
}

impl<'vm, 'jit: 'vm> Func<'jit, 'vm> {
    pub fn call(&self, vm: &'vm VM<'jit>, arg: Value<'jit, 'vm>) -> Result<Value<'jit, 'vm>> {
        use Param::*;

        let mut env = self.env;
        env = match self.func.param.clone() {
            Ident(ident) => env.enter_arg(ident, arg, &vm.bump),
            Formals {
                formals,
                ellipsis,
                alias,
            } => {
                let arg = arg.unwrap_attr_set();
                let mut new = HashMap::with_capacity_in(formals.len() + alias.iter().len(), &vm.bump);
                if !ellipsis
                    && arg
                        .as_inner()
                        .iter()
                        .map(|(k, _)| k)
                        .sorted()
                        .ne(formals.iter().map(|(k, _)| k).sorted())
                {
                    todo!()
                }
                for (formal, default) in formals {
                    let arg = arg
                        .select(formal)
                        .or_else(|| {
                            default.map(|idx| Value::Thunk(Thunk::new(vm.get_thunk(idx)).into()))
                        })
                        .unwrap();
                    new.insert(formal, arg);
                }
                if let Some(alias) = alias {
                    new.insert(alias, Value::AttrSet(arg));
                }
                env.enter_let(vm.bump.alloc(new), &vm.bump)
            }
        };

        let count = self.count.get();
        self.count.replace(count + 1);
        if count >= 1 {
            let compiled = self.compiled.get_or_init(|| vm.compile_func(self.func));
            let env = env as *const _;
            let ret = unsafe { compiled.call(vm as *const VM, env) };
            return Ok(ret.into());
        }
        vm.eval(self.func.opcodes.iter().copied(), vm.bump.alloc(env))
    }
}

impl PartialEq for Func<'_, '_> {
    fn eq(&self, _: &Self) -> bool {
        false
    }
}