use crate::{
    error::{Error, HasSpan},
    interpreter::{
        data::{CollectionKind, Data},
        generic_context::GenericContext,
        interpreter_conversion::ConversionKind,
        interpreter_generics::GenericValue,
        interpreter_run::Runnable,
        library::HardcodedId,
        namespace::{Namespace, NamespacedKind, Namespaces, ProgressKind, Typing},
        pointer::Pointer,
        scope::Scope,
        supported_type::SupportedType,
        Interpreter,
    },
};
use arcis_compiler::core::expressions::circuit::BaseCircuitId;
use quote::{format_ident, ToTokens};
use std::fmt::Debug;

#[derive(Debug, Clone)]
pub enum Function<'ast> {
    /// A function extracted from the AST of the circuit.
    /// Parameters are: signature, contents, outer scope.
    Ast(&'ast syn::Signature, &'ast syn::Block, Scope<'ast>),
    Circuit(BaseCircuitFunc<'ast>),
    Hardcoded(HardcodedFunc<'ast>),
    Static(&'static dyn FunctionTrait),
    StructRef(usize, Option<&'ast syn::Generics>),
}

impl<'ast> Function<'ast> {
    pub fn generics(&self) -> Option<&'ast syn::Generics> {
        match self {
            Function::Ast(sig, _, _) => Some(&sig.generics),
            Function::Circuit(func) => Some(&func.signature.generics),
            Function::Hardcoded(func) => Some(&func.signature.generics),
            Function::StructRef(_, a) => *a,
            _ => None,
        }
    }
    pub fn call_inner(
        &self,
        interpreter: &mut Interpreter<'ast>,
        args: &[Pointer],
        generics: Option<&'ast syn::Generics>,
        turbofish: Option<Vec<GenericValue>>,
        mut ctx: GenericContext,
        expr: &dyn ToTokens,
    ) -> Result<Pointer, Error> {
        interpreter.with_turbofish(
            turbofish,
            generics,
            &mut ctx,
            expr,
            |interpreter| match self {
                Function::Ast(sig, block, scope) => {
                    call_syn_function(sig, *block, *scope, interpreter, args)
                }
                Function::Circuit(func) => func.call(interpreter, args),
                Function::Hardcoded(func) => func.call(interpreter, args),
                Function::Static(func) => func.call(interpreter, args, expr),
                Function::StructRef(ty_id, _) => {
                    StructRef::new(*ty_id).call(interpreter, args, expr)
                }
            },
        )
    }
    pub fn call(
        &self,
        interpreter: &mut Interpreter<'ast>,
        args: &[Pointer],
        turbofish: Option<&'ast syn::AngleBracketedGenericArguments>,
        ctx: GenericContext,
        expr: &dyn ToTokens,
    ) -> Result<Pointer, Error> {
        let generics = self.generics();
        let turbofish = interpreter.evaluate_turbofish(turbofish, generics, expr)?;
        self.call_inner(interpreter, args, generics, turbofish, ctx, expr)
    }
    pub fn has_ref_receiver(&self) -> bool {
        fn has_sig_ref_receiver(sig: &syn::Signature) -> bool {
            sig.inputs.first().is_some_and(|t| {
                if let syn::FnArg::Receiver(r) = t {
                    return r.reference.is_some();
                };
                false
            })
        }
        match self {
            Function::Ast(sig, _, _) => has_sig_ref_receiver(sig),
            Function::Circuit(circ) => has_sig_ref_receiver(circ.signature),
            Function::Hardcoded(func) => has_sig_ref_receiver(func.signature),
            Function::Static(s) => s.has_ref_receiver(),
            Function::StructRef(..) => false,
        }
    }
}

pub trait FunctionTrait: Debug {
    fn call(
        &self,
        interpreter: &mut Interpreter,
        args: &[Pointer],
        expr: &dyn ToTokens,
    ) -> Result<Pointer, Error>;
    fn has_ref_receiver(&self) -> bool;
}

fn apply_sig_inputs<'ast>(
    sig: &'ast syn::Signature,
    interpreter: &mut Interpreter<'ast>,
    args: &[Pointer],
) -> Result<(), Error> {
    for (fn_arg, ptr) in sig.inputs.iter().zip(args) {
        let ptr = *ptr;
        let arg = interpreter.handle_fn_arg(fn_arg)?;
        match fn_arg {
            syn::FnArg::Receiver(_) => {
                let self_path = syn::parse_str("Self").expect("`Self` parsing failed.");
                let self_ident = format_ident!("Self");
                let Namespace::Type(ty_id) =
                    Namespaces::get_from_current_scope(interpreter, &self_ident)?
                else {
                    return Err(Error::new_spanned(
                        fn_arg,
                        "Function has a receiver but `Self` is unknown.",
                    ));
                };
                let (self_ty, None) =
                    Typing::get_from_index(interpreter, ty_id, Some(&self_path))?.destructure()
                else {
                    return Err(Error::new_spanned(fn_arg, "`Self` cannot have generics."));
                };
                let receiver_ty = if *arg.1 == syn::parse_str("Self").unwrap() {
                    self_ty.clone()
                } else if *arg.1 == syn::parse_str("&Self").unwrap()
                    || *arg.1 == syn::parse_str("&mut Self").unwrap()
                {
                    self_ty.reference()
                } else {
                    return Err(Error::new_spanned(
                        fn_arg,
                        "Only self, &self, &mut self are supported.",
                    ));
                };
                let ptr =
                    interpreter.conversion(ptr, &receiver_ty, ConversionKind::Implicit, fn_arg)?;
                let addr = interpreter.memory.clone_and_push(ptr)?;
                interpreter.new_binding_from_simple_address(format_ident!("self"), addr);
            }
            syn::FnArg::Typed(pat_ty) => interpreter.apply_pat_type(pat_ty, ptr, false)?,
        };
    }
    Ok(())
}

fn call_syn_function<'ast>(
    sig: &'ast syn::Signature,
    block: impl Runnable<'ast>,
    scope: Scope<'ast>,
    interpreter: &mut Interpreter<'ast>,
    args: &[Pointer],
) -> Result<Pointer, Error> {
    interpreter.with_scope(scope, |interpreter| {
        let output = interpreter.run_with_new_bindings(|interpreter| {
            apply_sig_inputs(sig, interpreter, args)?;
            block.run(interpreter)
        })?;
        let converted_output = if let syn::ReturnType::Type(_, return_ty) = &sig.output {
            let return_type = interpreter.type_from_syn_type(return_ty, None)?;
            interpreter.conversion(output, &return_type, ConversionKind::Implicit, return_ty)?
        } else {
            output
        };
        Ok(converted_output)
    })
}

#[derive(Debug)]
pub struct StructRef {
    ty_id: usize,
}

impl StructRef {
    pub fn new(ty_id: usize) -> Self {
        Self { ty_id }
    }
}
impl FunctionTrait for StructRef {
    fn call(
        &self,
        interpreter: &mut Interpreter,
        args: &[Pointer],
        expr: &dyn ToTokens,
    ) -> Result<Pointer, Error> {
        let ty_and_gen = Typing::get_from_index(interpreter, self.ty_id, None)?;
        let (SupportedType::Struct(mut struct_id), _) = ty_and_gen.destructure() else {
            return Err(Error::new_spanned(expr, "Not a struct"));
        };
        interpreter
            .refine_struct_type(&mut struct_id, false)
            .expect("Function failed with may_fail=false.");
        if args.len() != struct_id.get_n_fields() {
            return Err(Error::new_spanned(expr, "wrong number of arguments."));
        }
        let addresses: Result<Vec<_>, _> = args
            .iter()
            .zip(struct_id.types())
            .map(|(ptr, ty)| {
                let ptr = *ptr;
                let converted_ptr =
                    interpreter.conversion(ptr, ty, ConversionKind::Implicit, expr)?;
                interpreter.memory.clone_and_push(converted_ptr)
            })
            .collect();
        let addresses = addresses?;
        interpreter
            .refine_struct_type(&mut struct_id, true)
            .map_err(|e| {
                Error::new_spanned(
                    expr,
                    format!(
                        "Failed to define struct type: {e} in {}.",
                        expr.to_token_stream()
                    ),
                )
            })?;
        Ok(interpreter
            .memory
            .push(Data::Collection(
                addresses,
                CollectionKind::Struct(struct_id.clone()),
            ))
            .as_pointer())
    }
    fn has_ref_receiver(&self) -> bool {
        false
    }
}

#[derive(Debug, Clone, Copy)]
pub struct BaseCircuitFunc<'ast> {
    circuit: BaseCircuitId,
    signature: &'ast syn::Signature,
    outer_scope: Scope<'ast>,
}

impl<'ast> BaseCircuitFunc<'ast> {
    pub fn new(
        circuit: BaseCircuitId,
        signature: &'ast syn::Signature,
        outer_scope: Scope<'ast>,
    ) -> Self {
        BaseCircuitFunc {
            circuit,
            signature,
            outer_scope,
        }
    }
    fn call(
        &self,
        interpreter: &mut Interpreter<'ast>,
        args: &[Pointer],
    ) -> Result<Pointer, Error> {
        interpreter.with_scope(self.outer_scope, |interpreter| {
            interpreter.run_with_new_bindings(|interpreter| {
                apply_sig_inputs(self.signature, interpreter, args)
            })
        })?;
        let syn::ReturnType::Type(_, syn_return_ty) = &self.signature.output else {
            return Err(Error::new_spanned(
                self.signature,
                "Circuit functions should have a return type.",
            ));
        };
        let mut input_values = Vec::new();
        for arg in args {
            let data = arg.as_data();
            interpreter.data_to_arcis_values_impl(data, &mut input_values)?;
        }
        let output_values = self.circuit.apply(input_values);
        let mut return_ty = interpreter.with_scope(self.outer_scope, |interpreter| {
            interpreter.type_from_syn_type(syn_return_ty, None)
        })?;
        interpreter.refine_type(&mut return_ty, true).map_err(|e| {
            Error::new_spanned(syn_return_ty, format!("Not an explicit type: {:?}", e))
        })?;
        let expected_output_values = return_ty
            .n_arcis_values()
            .expect("The type should have just been made explicit.");
        if output_values.len() != expected_output_values {
            return Err(Error::new_spanned(
                syn_return_ty,
                format!(
                    "Circuit returns {} values, expected {}.",
                    output_values.len(),
                    expected_output_values
                ),
            ));
        }
        let result = interpreter.build_from_arcis_values(
            &mut output_values.as_slice(),
            &return_ty,
            false,
            self.signature.span(),
        )?;
        Ok(result.as_pointer())
    }
}

#[derive(Debug, Clone, Copy)]
pub struct HardcodedFunc<'ast> {
    func: HardcodedId,
    signature: &'ast syn::Signature,
    outer_scope: Scope<'ast>,
}

impl<'ast> HardcodedFunc<'ast> {
    pub fn new(
        func: HardcodedId,
        signature: &'ast syn::Signature,
        outer_scope: Scope<'ast>,
    ) -> Self {
        HardcodedFunc {
            func,
            signature,
            outer_scope,
        }
    }
    pub fn call(
        &self,
        interpreter: &mut Interpreter<'ast>,
        args: &[Pointer],
    ) -> Result<Pointer, Error> {
        call_syn_function(
            self.signature,
            (self.func, args, self.signature),
            self.outer_scope,
            interpreter,
            args,
        )
    }
}
