use crate::{
    error::{Error, HasSpan},
    interpreter::{
        data::{ArcisField, ArcisValue, CollectionKind, Data, ValData},
        namespace::{Namespace, NamespacedKind, Namespaces, ProgressKind, Typing},
        pointer::Pointer,
        supported_type::{ArrayLen, IntegerInfo, RustType, SupportedType},
        Interpreter,
    },
};
use arcis_compiler::{
    traits::{FromLeBits, GetBit},
    utils::{
        packing::{DataSize, PackLocation},
        used_field::UsedField,
    },
};
use ff::Field;
use quote::{format_ident, ToTokens};

fn gather_data_sizes_and_signedness(
    supported_type: &SupportedType,
    data_sizes: &mut Vec<DataSize>,
    signedness: &mut Vec<bool>,
) -> Result<(), &'static str> {
    match supported_type {
        SupportedType::Unit => {}
        SupportedType::FieldValue(r) => {
            data_sizes.push(r.data_size()?);
            signedness.push(r.is_signed());
        }
        SupportedType::Bool => {
            data_sizes.push(DataSize::PowerOfTwo(0));
            signedness.push(false)
        }
        SupportedType::Curve(_) => {
            data_sizes.push(DataSize::Full);
            signedness.push(false)
        }
        SupportedType::Struct(s) => {
            for ty in s.types() {
                gather_data_sizes_and_signedness(ty, data_sizes, signedness)?;
            }
        }
        SupportedType::Ref(ty) => {
            gather_data_sizes_and_signedness(ty, data_sizes, signedness)?;
        }
        SupportedType::Array(ty, len) => {
            let len = match len {
                ArrayLen::Computed(len) => *len,
                ArrayLen::ConstVar(_) => return Err("Cannot compute generic length."),
            };
            for _ in 0..len {
                gather_data_sizes_and_signedness(ty, data_sizes, signedness)?;
            }
        }
        SupportedType::Tuple(ty) => {
            for ty in ty {
                gather_data_sizes_and_signedness(ty, data_sizes, signedness)?;
            }
        }
        SupportedType::Slice(ty, len) => {
            let len = len.value()?;
            gather_data_sizes_and_signedness(
                &SupportedType::Array(ty.clone(), ArrayLen::Computed(len)),
                data_sizes,
                signedness,
            )?;
        }
        _ => return Err("Unsupported type"),
    };
    Ok(())
}

type FullPackLocation = (PackLocation, DataSize, bool);
pub fn packing(
    supported_type: &SupportedType,
) -> Result<(Vec<FullPackLocation>, Vec<Option<u16>>), &'static str> {
    let mut data_sizes = Vec::new();
    let mut signedness = Vec::new();
    gather_data_sizes_and_signedness(supported_type, &mut data_sizes, &mut signedness)?;
    let len = data_sizes.len();
    let (packing, container_states) = DataSize::pack_arcis(data_sizes.clone());
    let res = (0..len)
        .map(|idx| (packing[idx], data_sizes[idx], signedness[idx]))
        .collect();
    Ok((res, container_states))
}

fn unpack(
    in_values: &[ArcisValue],
    packing: &[FullPackLocation],
    container_states: &[Option<u16>],
) -> Vec<ArcisValue> {
    assert_eq!(in_values.len(), container_states.len());
    let in_values = in_values
        .iter()
        .copied()
        .enumerate()
        .map(|(idx, v)| {
            if let Some(width) = container_states[idx] {
                RustType::Integer(Some(IntegerInfo {
                    signed: false,
                    width: width as usize,
                }))
                .cap_arcis_value(v)
            } else {
                v
            }
        })
        .collect::<Vec<_>>();
    packing
        .iter()
        .map(|(pack_location, size, signedness)| match size {
            DataSize::Full => {
                assert_eq!(pack_location.bit_offset, 0);
                assert!(!signedness);
                in_values[pack_location.index]
            }
            DataSize::PowerOfTwo(n) => {
                let in_value = in_values[pack_location.index];
                let begin = pack_location.bit_offset;
                let end = begin + (1u16 << n);
                ArcisValue::from_le_bits(
                    (begin..end)
                        .map(|i| {
                            let bit = in_value.get_bit(i as usize, false);
                            if i == end - 1 && *signedness {
                                !bit
                            } else {
                                bit
                            }
                        })
                        .collect(),
                    *signedness,
                )
            }
        })
        .collect()
}
fn pack(
    in_values: &[ArcisValue],
    packing: &[FullPackLocation],
    packing_size: usize,
) -> Vec<ArcisValue> {
    assert_eq!(in_values.len(), packing.len());
    let mut result = vec![ArcisValue::from(0); packing_size];
    for (val, (location, size, signedness)) in std::iter::zip(in_values, packing) {
        let offset = match size {
            DataSize::PowerOfTwo(exp) => {
                if *signedness {
                    ArcisField::power_of_two((1usize << exp) - 1)
                } else {
                    ArcisField::ZERO
                }
            }
            DataSize::Full => {
                assert_eq!(location.bit_offset, 0);
                assert!(!*signedness);
                ArcisField::ZERO
            }
        };
        result[location.index] +=
            (*val + offset) * ArcisField::power_of_two(location.bit_offset as usize);
    }
    result
}

pub fn pack_func(
    interpreter: &mut Interpreter,
    args: &[Pointer],
    expr: &dyn ToTokens,
) -> Result<Pointer, Error> {
    if args.len() != 1 {
        return Err(Error::new_spanned(expr, "Function should have 1 argument."));
    }
    let arg = args[0];
    let vals = interpreter.pointer_to_arcis_values(arg)?;
    let type_ident = format_ident!("T");
    let namespace = Namespaces::get_from_current_scope(interpreter, &type_ident)?;
    let Namespace::Type(ty_id) = namespace else {
        return Err(Error::new_spanned(expr, "Argument should be a type."));
    };
    let ty_and_gen = Typing::get_from_index(interpreter, ty_id, None)?;

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

    let (pack_locations, pack_loads) = packing(&ty).map_err(|err| {
        Error::new_spanned(expr, format!("Cannot pack type `{}` because {}.", ty, err))
    })?;
    let packed_vals = pack(&vals, &pack_locations, pack_loads.len());
    interpreter.client_types.insert(ty).map_err(|err| {
        Error::new_spanned(
            expr,
            format!("Problem when trying to tell client how to pack: {}", err),
        )
    })?;

    let rust_type = RustType::ArcisField;
    let addrs = packed_vals
        .into_iter()
        .map(|val| interpreter.memory.push(Data::Val(ValData(val, rust_type))))
        .collect();
    Ok(interpreter
        .memory
        .push(Data::Collection(
            addrs,
            CollectionKind::Array(SupportedType::FieldValue(rust_type)),
        ))
        .as_pointer())
}

pub fn unpack_func(
    interpreter: &mut Interpreter,
    args: &[Pointer],
    expr: &dyn ToTokens,
) -> Result<Pointer, Error> {
    if args.len() != 1 {
        return Err(Error::new_spanned(expr, "Function should have 1 argument."));
    }
    let arg = args[0];
    let vals = interpreter.pointer_to_arcis_values(arg)?;
    let type_ident = format_ident!("T");
    let namespace = Namespaces::get_from_current_scope(interpreter, &type_ident)?;
    let Namespace::Type(ty_id) = namespace else {
        return Err(Error::new_spanned(expr, "Argument should be a type."));
    };
    let ty_and_gen = Typing::get_from_index(interpreter, ty_id, None)?;

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

    let (pack_locations, pack_loads) = packing(&ty).map_err(|err| {
        Error::new_spanned(
            expr,
            format!("Cannot unpack type `{}` because {}.", ty, err),
        )
    })?;
    let unpacked_vals = unpack(&vals, &pack_locations, &pack_loads);
    let res = interpreter.build_from_arcis_values(
        &mut unpacked_vals.as_slice(),
        &ty,
        false,
        expr.span(),
    )?;
    interpreter.client_types.insert(ty).map_err(|err| {
        Error::new_spanned(
            expr,
            format!("Problem when trying to tell client how to pack: {}", err),
        )
    })?;
    Ok(res.as_pointer())
}
