use crate::{
    error::Error,
    interpreter::{
        data::{BoolData, Data},
        data_store::SimpleAddress,
        generic_context::GenericContext,
        interpreter_conversion::ConversionKind,
        interpreter_generics::GenericValue,
        interpreter_statics::Functions,
        namespace::ProgressKind,
        pointer::Pointer,
        scope::Scope,
        Interpreter,
    },
    utils::Ptr,
};
use arcis_compiler::core::circuits::boolean::boolean_value::BooleanValue;
use quote::ToTokens;
use std::{cell::Cell, collections::HashMap};
use syn::ReturnType;

#[derive(Debug, Clone, PartialEq)]
pub struct ConditionalFn {
    cond: BoolData,
    then_fn_address: SimpleAddress,
    else_fn_address: SimpleAddress,
}

impl ConditionalFn {
    pub fn new(
        cond: BooleanValue,
        then_fn_address: SimpleAddress,
        else_fn_address: SimpleAddress,
    ) -> Self {
        let cond = BoolData(cond);
        Self {
            cond,
            then_fn_address,
            else_fn_address,
        }
    }
}
#[derive(Debug, Clone, PartialEq)]
pub struct IndexFn {
    ohe_address: SimpleAddress,
    fn_arr_address: SimpleAddress,
}

impl IndexFn {
    pub fn new(ohe_address: SimpleAddress, fn_arr_address: SimpleAddress) -> Self {
        Self {
            ohe_address,
            fn_arr_address,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct Closure<'ast> {
    expr: Ptr<'ast, syn::ExprClosure>,
    scope: Scope<'ast>,
}

impl<'ast> Closure<'ast> {
    pub fn new(expr: &'ast syn::ExprClosure, scope: Scope<'ast>) -> Self {
        let expr = Ptr(expr);
        Self { expr, scope }
    }
    pub fn call(
        &self,
        interpreter: &mut Interpreter<'ast>,
        args: &[Pointer],
        expr: &dyn ToTokens,
        new_bindings: &HashMap<syn::Ident, SimpleAddress>,
    ) -> Result<Pointer, Error> {
        let closure_expr = self.expr.0;
        interpreter.with_scope(self.scope, |interpreter| {
            interpreter.run_with_new_bindings(|interpreter| {
                interpreter.bindings = new_bindings.clone();
                if closure_expr.inputs.len() != args.len() {
                    return Err(Error::new_spanned(
                        expr,
                        "Wrong number of arguments in closure call.",
                    ));
                }
                for (pat, arg) in closure_expr.inputs.iter().zip(args) {
                    interpreter.apply_pat(pat, *arg, false)?;
                }
                let res = interpreter.run_expr(&closure_expr.body)?;
                let converted_res = match &closure_expr.output {
                    ReturnType::Default => res,
                    ReturnType::Type(_, ty) => {
                        let supported_ty = interpreter.type_from_syn_type(ty, None)?;
                        interpreter.conversion(res, &supported_ty, ConversionKind::Implicit, ty)?
                    }
                };
                Ok(converted_res)
            })
        })
    }
}

/// Values of function variables.
#[derive(Debug, Clone, PartialEq)]
pub enum FnData {
    /// A `Function` stored in the `Interpreter`.
    Simple(usize, Option<Vec<GenericValue>>, GenericContext),
    /// A closure.
    Closure(usize, HashMap<syn::Ident, SimpleAddress>),
    /// Used for `if c { fn_a } else { fn_b }`.
    Cond(ConditionalFn),
    /// Used for `array_of_functions[idx]`.
    Index(IndexFn),
}

impl FnData {
    pub fn call(
        &self,
        interpreter: &mut Interpreter,
        args: &[Pointer],
        expr: &dyn ToTokens,
    ) -> Result<Pointer, Error> {
        match self.clone() {
            FnData::Simple(a, gen, ctx) => {
                let func = Functions::get_from_index(interpreter, a, None)?;
                let generics = func.generics();
                func.call_inner(interpreter, args, generics, gen, ctx, expr)
            }
            FnData::Cond(func) => interpreter.conditionally_run(
                func.cond.0,
                (func.then_fn_address, args, expr),
                Some((func.else_fn_address, args, expr)),
            ),
            FnData::Index(func) => {
                let Data::Collection(ohe, _) = interpreter.memory[func.ohe_address].clone() else {
                    return Err(Error::new_spanned(expr, "OHE array is not an array."));
                };
                let Data::Collection(functions, _) =
                    interpreter.memory[func.fn_arr_address].clone()
                else {
                    return Err(Error::new_spanned(expr, "OHE array is not an array."));
                };
                if ohe.len() != functions.len() {
                    return Err(Error::new_spanned(
                        expr,
                        "Length of OHE and array do not match.",
                    ));
                }
                let mut pointers = Vec::with_capacity(functions.len());
                let mut ohe_vec = Vec::with_capacity(functions.len());
                for (ohe, function) in ohe.into_iter().zip(functions) {
                    let Data::Bool(BoolData(ohe)) = interpreter.memory[ohe] else {
                        return Err(Error::new_spanned(expr, "OHE array is not an OHE array."));
                    };
                    ohe_vec.push(ohe);
                    let none_pointer = Option::<Pointer>::None;
                    let res = Cell::new(none_pointer);
                    interpreter.conditionally_run(
                        ohe,
                        ((function, args, expr), &res),
                        none_pointer,
                    )?;
                    let res = res.take().expect("Should have ran.");
                    pointers.push(res);
                }
                Ok(interpreter
                    .memory
                    .merge_array_of_pointers(pointers[0], &ohe_vec, &pointers))
            }
            FnData::Closure(id, bindings) => {
                let closure = interpreter.closures[id];
                closure.call(interpreter, args, expr, &bindings)
            }
        }
    }
}
