use crate::{
    error::Error,
    interpreter::{
        data::{BoolData, CollectionKind, Data, ValData},
        function::FunctionTrait,
        interpreter_run::Runnable,
        library::{
            cipher::{
                decrypt_func,
                encrypt_client_func,
                encrypt_mxe_func,
                x25519_pubkey_from_x,
                x25519_pubkey_to_x,
            },
            packing::{pack_func, packing, unpack_func},
            sort,
        },
        namespace::{Namespace, NamespacedKind, Namespaces, ProgressKind, Typing},
        pointer::Pointer,
        supported_type::{RustType, SupportedType},
        Interpreter,
    },
};
use quote::{format_ident, ToTokens};
use std::fmt::Debug;

#[derive(Debug, Clone, Copy)]
pub enum HardcodedId {
    Len,
    Clone,
    Reveal,
    IntoIter,
    NArcisValues,
    ToArcis,
    MxeEncrypt,
    ClientEncrypt,
    NBools,
    FromBools,
    X25519PubkeyFromX,
    X25519PubkeyToX,
    NPackedArcisValues,
    Pack,
    Unpack,
    ArcisRNGShuffle,
    ArcisRNGBool,
    ArcisRNGGenIntegerFromWidth,
    ArcisRNGGenPublicIntegerFromWidth,
}

impl HardcodedId {
    pub fn new_from_str(ident: &str) -> Option<Self> {
        match ident {
            "len" => Some(HardcodedId::Len),
            "clone" => Some(HardcodedId::Clone),
            "reveal" => Some(HardcodedId::Reveal),
            "into_iter" => Some(HardcodedId::IntoIter),
            "n_arcis_values" => Some(HardcodedId::NArcisValues),
            "to_arcis" => Some(HardcodedId::ToArcis),
            "mxe_encrypt" => Some(HardcodedId::MxeEncrypt),
            "client_encrypt" => Some(HardcodedId::ClientEncrypt),
            "n_bools" => Some(HardcodedId::NBools),
            "from_bools" => Some(HardcodedId::FromBools),
            "x25519_pubkey_from_x" => Some(HardcodedId::X25519PubkeyFromX),
            "x25519_pubkey_to_x" => Some(HardcodedId::X25519PubkeyToX),
            "n_packed_arcis_values" => Some(HardcodedId::NPackedArcisValues),
            "pack" => Some(HardcodedId::Pack),
            "unpack" => Some(HardcodedId::Unpack),
            "ArcisRNG::shuffle" => Some(HardcodedId::ArcisRNGShuffle),
            "ArcisRNG::bool" => Some(HardcodedId::ArcisRNGBool),
            "ArcisRNG::gen_integer_from_width" => Some(HardcodedId::ArcisRNGGenIntegerFromWidth),
            "ArcisRNG::gen_public_integer_from_width" => {
                Some(HardcodedId::ArcisRNGGenPublicIntegerFromWidth)
            }
            _ => None,
        }
    }
    fn call(
        &self,
        interpreter: &mut Interpreter,
        args: &[Pointer],
        expr: &dyn ToTokens,
    ) -> Result<Pointer, Error> {
        let res = (match self {
            HardcodedId::Len => len,
            HardcodedId::Clone => clone,
            HardcodedId::Reveal => reveal,
            HardcodedId::IntoIter => into_iter,
            HardcodedId::NArcisValues => n_arcis_values,
            HardcodedId::ToArcis => decrypt_func,
            HardcodedId::MxeEncrypt => encrypt_mxe_func,
            HardcodedId::ClientEncrypt => encrypt_client_func,
            HardcodedId::NBools => n_bools,
            HardcodedId::FromBools => from_bools,
            HardcodedId::X25519PubkeyFromX => x25519_pubkey_from_x,
            HardcodedId::X25519PubkeyToX => x25519_pubkey_to_x,
            HardcodedId::NPackedArcisValues => n_packed_arcis_values,
            HardcodedId::Pack => pack_func,
            HardcodedId::Unpack => unpack_func,
            HardcodedId::ArcisRNGShuffle => sort::shuffle,
            HardcodedId::ArcisRNGBool => sort::gen_bool,
            HardcodedId::ArcisRNGGenIntegerFromWidth => sort::gen_integer_from_width,
            HardcodedId::ArcisRNGGenPublicIntegerFromWidth => sort::gen_public_integer_from_width,
        })(interpreter, args, expr)?;

        Ok(res)
    }
}

impl<'ast, T: ToTokens + Copy> Runnable<'ast> for (HardcodedId, &[Pointer], T) {
    fn run(self, interpreter: &mut Interpreter<'ast>) -> Result<Pointer, Error> {
        let (h, args, tokens) = self;
        h.call(interpreter, args, &tokens)
    }
}

#[derive(Clone, Copy)]
pub struct Hardcoded {
    inner: fn(
        interpreter: &mut Interpreter,
        args: &[Pointer],
        expr: &dyn ToTokens,
    ) -> Result<Pointer, Error>,
    name: &'static str,
    has_ref_receiver: bool,
}

impl Debug for Hardcoded {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str("StdFunc(")
            .and_then(|_| f.write_str(self.name))
            .and_then(|_| f.write_str(")"))
    }
}

impl FunctionTrait for Hardcoded {
    fn call(
        &self,
        interpreter: &mut Interpreter,
        args: &[Pointer],
        expr: &dyn ToTokens,
    ) -> Result<Pointer, Error> {
        (self.inner)(interpreter, args, expr)
    }
    fn has_ref_receiver(&self) -> bool {
        self.has_ref_receiver
    }
}

impl Hardcoded {
    pub const fn new(
        name: &'static str,
        inner: fn(
            interpreter: &mut Interpreter,
            args: &[Pointer],
            expr: &dyn ToTokens,
        ) -> Result<Pointer, Error>,
        has_ref_receiver: bool,
    ) -> Self {
        Hardcoded {
            inner,
            name,
            has_ref_receiver,
        }
    }
}

pub fn len(
    interpreter: &mut Interpreter,
    args: &[Pointer],
    expr: &dyn ToTokens,
) -> Result<Pointer, Error> {
    let arg = args[0];
    let data = interpreter.memory.get_underlying_data(arg)?;
    let Data::Collection(v, _) = data else {
        return Err(Error::new_spanned(expr, "unexpected data"));
    };
    let len = v.len();

    Ok(interpreter
        .memory
        .push(Data::Val(ValData(len.into(), RustType::usize())))
        .as_pointer())
}

pub fn clone(
    interpreter: &mut Interpreter,
    args: &[Pointer],
    expr: &dyn ToTokens,
) -> Result<Pointer, Error> {
    let arg = args[0];
    let data = interpreter.memory.compute(arg)?;
    let Data::Pointer(ptr) = data else {
        return Err(Error::new_spanned(expr, "unexpected data"));
    };
    let res_addr = interpreter.memory.clone_and_push(ptr)?;
    Ok(res_addr.as_pointer())
}

pub fn reveal(
    interpreter: &mut Interpreter,
    args: &[Pointer],
    expr: &dyn ToTokens,
) -> Result<Pointer, Error> {
    if interpreter.memory.in_conditional_execution() {
        return Err(interpreter
            .backtrace
            .make_conditional_execution_error(expr, ".reveal()"));
    }
    let arg = args[0];
    let data = interpreter.memory.compute(arg)?;
    let revealed_data = data.reveal(&mut interpreter.memory)?;
    Ok(interpreter.memory.push(revealed_data).as_pointer())
}

pub fn into_iter(
    interpreter: &mut Interpreter,
    args: &[Pointer],
    expr: &dyn ToTokens,
) -> Result<Pointer, Error> {
    let arg = args[0];
    let receiver_type = interpreter.memory.pointed_type(arg);
    let is_ref = receiver_type.is_reference();
    let is_slice = receiver_type.underlying_type().is_slice();
    if is_slice {
        if is_ref {
            interpreter.run_method(&format_ident!("iter"), &mut [arg], None, expr, None)
        } else {
            interpreter.run_method(&format_ident!("__into_iter"), &mut [arg], None, expr, None)
        }
    } else {
        Ok(arg)
    }
}

fn n_arcis_values(
    interpreter: &mut Interpreter,
    args: &[Pointer],
    expr: &dyn ToTokens,
) -> Result<Pointer, Error> {
    if !args.is_empty() {
        return Err(Error::new_spanned(
            expr,
            "Function should have no arguments.",
        ));
    }
    let self_ident = format_ident!("T");
    let namespace = Namespaces::get_from_current_scope(interpreter, &self_ident)?;
    let Namespace::Type(ty_id) = namespace else {
        return Err(Error::new_spanned(expr, "Argument should be a type."));
    };
    let self_ty_and_gen = Typing::get_from_index(interpreter, ty_id, None)?;

    let (mut self_ty, None) = self_ty_and_gen.destructure() else {
        return Err(Error::new_spanned(expr, "Generics on `Self`."));
    };
    interpreter.refine_type(&mut self_ty, true).map_err(|err| {
        Error::new_spanned(
            expr,
            format!("Cannot count type {self_ty:?} because {}.", err),
        )
    })?;

    let res = self_ty.n_arcis_values().expect("We just refined the type.");

    Ok(interpreter
        .memory
        .push(Data::Val(ValData(res.into(), RustType::usize())))
        .as_pointer())
}

fn n_packed_arcis_values(
    interpreter: &mut Interpreter,
    args: &[Pointer],
    expr: &dyn ToTokens,
) -> Result<Pointer, Error> {
    if !args.is_empty() {
        return Err(Error::new_spanned(
            expr,
            "Function should have no arguments.",
        ));
    }
    let self_ident = format_ident!("T");
    let namespace = Namespaces::get_from_current_scope(interpreter, &self_ident)?;

    let Namespace::Type(ty_id) = namespace else {
        return Err(Error::new_spanned(expr, "Argument should be a type."));
    };
    let self_ty_and_gen = Typing::get_from_index(interpreter, ty_id, None)?;

    let (mut self_ty, None) = self_ty_and_gen.destructure() else {
        return Err(Error::new_spanned(expr, "Generics on `Self`."));
    };
    interpreter.refine_type(&mut self_ty, true).map_err(|err| {
        Error::new_spanned(
            expr,
            format!("Cannot count type {self_ty:?} because {}.", err),
        )
    })?;

    let res = packing(&self_ty)
        .expect("We just refined the type.")
        .1
        .len();

    Ok(interpreter
        .memory
        .push(Data::Val(ValData(res.into(), RustType::usize())))
        .as_pointer())
}

fn n_bools(
    interpreter: &mut Interpreter,
    args: &[Pointer],
    expr: &dyn ToTokens,
) -> Result<Pointer, Error> {
    if !args.is_empty() {
        return Err(Error::new_spanned(
            expr,
            "Function should have no arguments.",
        ));
    }
    let self_ident = format_ident!("T");
    let namespace = Namespaces::get_from_current_scope(interpreter, &self_ident)?;
    let Namespace::Type(ty_id) = namespace else {
        return Err(Error::new_spanned(expr, "Argument should be a type."));
    };
    let self_ty_and_gen = Typing::get_from_index(interpreter, ty_id, None)?;

    let (mut self_ty, None) = self_ty_and_gen.destructure() else {
        return Err(Error::new_spanned(expr, "Generics on `Self`."));
    };
    interpreter.refine_type(&mut self_ty, true).map_err(|err| {
        Error::new_spanned(
            expr,
            format!("Cannot count type {self_ty:?} because {}.", err),
        )
    })?;

    let res = self_ty.n_bools().map_err(|err| {
        Error::new_spanned(
            expr,
            format!("Cannot count type {self_ty:?} because {}.", err),
        )
    })?;

    Ok(interpreter
        .memory
        .push(Data::Val(ValData(res.into(), RustType::usize())))
        .as_pointer())
}

fn from_bools(
    interpreter: &mut Interpreter,
    args: &[Pointer],
    expr: &dyn ToTokens,
) -> Result<Pointer, Error> {
    let self_ident = format_ident!("T");
    let namespace = Namespaces::get_from_current_scope(interpreter, &self_ident)?;
    let Namespace::Type(ty_id) = namespace else {
        return Err(Error::new_spanned(
            expr,
            "Generic argument should be a type.",
        ));
    };
    let self_ty_and_gen = Typing::get_from_index(interpreter, ty_id, None)?;

    let (mut self_ty, None) = self_ty_and_gen.destructure() else {
        return Err(Error::new_spanned(expr, "Generics on `Self`."));
    };
    interpreter.refine_type(&mut self_ty, true).map_err(|err| {
        Error::new_spanned(
            expr,
            format!("Cannot count type {self_ty:?} because {}.", err),
        )
    })?;

    let n_bools = self_ty.n_bools().map_err(|err| {
        Error::new_spanned(
            expr,
            format!("Cannot count type {self_ty:?} because {}.", err),
        )
    })?;

    let arg = args[0];
    let Data::Collection(v, CollectionKind::Array(SupportedType::Bool)) =
        interpreter.memory.get_underlying_data(arg)?
    else {
        return Err(Error::new_spanned(
            expr,
            "Argument should be an array of booleans.",
        ));
    };
    if v.len() != n_bools {
        return Err(Error::new_spanned(
            expr,
            "Wrong number of items in array of booleans.",
        ));
    }
    let mut data = Vec::with_capacity(n_bools);
    for addr in v {
        let Data::Bool(BoolData(val)) = interpreter.memory[addr] else {
            return Err(Error::new_spanned(
                expr,
                "Argument should be an array of booleans.",
            ));
        };
        data.push(val);
    }
    data.reverse();
    let res = self_ty.build_from_bools(&mut data, &mut interpreter.memory);

    Ok(res.as_pointer())
}
