use std::cell::OnceCell;
use core::mem::MaybeUninit;
use std::rc::Rc;

use hashbrown::{HashMap, HashSet};
use priority_queue::PriorityQueue;

use crate::env::Env;
use crate::error::Result;
use crate::eval::jit::{JITCompiler, JITFunc};
use crate::eval::Evaluate;
use crate::ir::{Dep, Downgraded, Ir, SccNode};
use crate::ty::internal as i;
use crate::ty::public::Value;

#[cfg(test)]
mod test;

type ThunkIdx = usize;
type EnvIdx = usize;

pub struct Engine {
    pub thunks: Box<[Ir]>,
    pub funcs: Box<[Ir]>,
    pub func_deps: Vec<HashMap<Dep, usize>>,
    jit: JITCompiler,
    compiled: Box<[OnceCell<JITFunc>]>,
    tasks: PriorityQueue<CompileTask, usize>,
}

pub fn eval(downgraded: Downgraded) -> Result<Value> {
    let mut engine = Engine::new(
        downgraded.thunks,
        downgraded.funcs,
        downgraded.func_deps,
        JITCompiler::new()
    );
    engine.eval(downgraded.graph)
}

impl Engine {
    pub fn new(thunks: Box<[Ir]>, funcs: Box<[Ir]>, func_deps: Vec<HashMap<Dep, usize>>, jit: JITCompiler) -> Self {
        Self {
            compiled: (0..thunks.len() + funcs.len()).map(|_| OnceCell::new()).collect(),
            tasks: PriorityQueue::new(),
            thunks,
            funcs,
            func_deps,
            jit,
        }
    }

    pub fn eval(&mut self, graph: Vec<SccNode>) -> Result<Value> {
        let mut env = Env::new();
        let last = graph.last().unwrap().members[0];
        for SccNode { members, .. } in graph.into_iter() {
            if members.len() == 1 {
                for member in members.into_iter() {
                    let engine = unsafe { &mut *(self as *mut Self) };
                    let mut val = self.thunks[member].eval(engine, &mut env)?;
                    val.force(engine, &mut env)?;
                    env.insert_cache(member, val);
                }
            } else {
                todo!();
                for member in members.into_iter() {
                    let val = self.thunks[member].clone().eval(self, &mut env)?;
                    env.insert_cache(member, val);
                }
            }
        }
        env.lookup_cache(last, |_| unreachable!()).map(|mut val| {
            Ok(val
                .force(self, &mut env)?
                .to_public(self, &mut HashSet::new()))
        })?
    }

    pub fn eval_thunk(&mut self, idx: usize, env: &mut Env) -> Result<i::Value> {
        let engine = self as *const Engine;
        let func =  self.compiled[idx].get_or_init(|| self.jit.compile(&self.thunks[idx], idx));
        let mut ret: MaybeUninit<i::Value> = MaybeUninit::uninit();
        unsafe {
            func(engine, env, ret.as_mut_ptr());
            Ok(ret.assume_init())
        }
    }

    pub fn call_func(&mut self, idx: usize, env: &mut Env) -> Result<i::Value> {
        let engine = self as *const Engine;
        let func = self.compiled[idx + self.thunks.len()].get_or_init(|| self.jit.compile(&self.funcs[idx], idx));
        let mut ret: MaybeUninit<i::Value> = MaybeUninit::uninit();
        unsafe {
            func(engine, env, ret.as_mut_ptr());
            Ok(ret.assume_init())
        }
    }

    pub fn eval_func_deps(&mut self, idx: usize, env: &mut Env) -> Result<()> {
        for (&dep, _) in
            unsafe { &*(&self.func_deps[idx] as *const HashMap<Dep, usize>) }.iter()
        {
            match dep {
                Dep::Arg(idx) => {
                    if let i::Value::Thunk(idx) = env.lookup_arg(idx) {
                        env.insert_cache_lazy(idx, |env| {
                            let engine = unsafe { &mut *(self as *mut Self) };
                            let mut val = self.thunks[idx].eval(engine, env)?;
                            val.force(engine, env)?;
                            val.ok()
                        })?;
                    }
                }
                Dep::Thunk(idx) => {
                    env.insert_cache_lazy(idx, |env| {
                        let engine = unsafe { &mut *(self as *mut Self) };
                        let mut val = self.thunks[idx].eval(engine, env)?;
                        val.force(engine, env)?;
                        val.ok()
                    })?;
                }
            }
        }
        Ok(())
    }
}

enum Thunk {
    Expr(Ir),
    Compiling,
    Compiled(fn(Rc<Env>) -> Value),
}

impl Thunk {
    fn new(ir: Ir) -> Thunk {
        Thunk::Expr(ir)
    }
}

#[derive(Hash, PartialEq, Eq)]
struct CompileTask {
    idx: usize,
}