use crate::interpreter::{
    data::{ArcisField, BoolData, CollectionKind, CurveData, Data},
    data_store::SimpleAddress,
    memory::Memory,
    supported_type::{ArrayLen, SupportedType},
};
use arcis_compiler::core::{
    expressions::{
        curve_expr,
        curve_expr::CurveExpr,
        expr::Expr,
        field_expr,
        field_expr::{FieldExpr, InputId},
        InputKind,
    },
    global_value::{curve_value::CurveValue, value::FieldValue},
};
use ff::Field;
use std::rc::Rc;

pub struct InputsHandlingState {
    input_id: InputId,
}
impl InputsHandlingState {
    pub fn new() -> Self {
        InputsHandlingState { input_id: 0 }
    }
    pub fn create_inputs(
        &mut self,
        memory: &mut Memory,
        supported_type: SupportedType,
        name: String,
    ) -> Result<SimpleAddress, &'static str> {
        let address = match supported_type {
            SupportedType::Unit => Memory::SIMPLE_ADDRESS_OF_UNIT,
            SupportedType::FieldValue(r) => {
                let (min, max) = r.min_and_max();
                let expr = FieldExpr::Input(
                    self.input_id,
                    Rc::new(field_expr::InputInfo::<ArcisField> {
                        kind: InputKind::Plaintext,
                        min,
                        max,
                        name,
                        ..Default::default()
                    }),
                );
                self.input_id += 1;
                let field_value = FieldValue::new(expr);
                memory.push(Data::Val(r.build_from_arcis_value(field_value)))
            }
            SupportedType::Bool => {
                let (min, max) = (ArcisField::ZERO, ArcisField::ONE);
                let expr = FieldExpr::Input(
                    self.input_id,
                    Rc::new(field_expr::InputInfo::<ArcisField> {
                        kind: InputKind::Plaintext,
                        min,
                        max,
                        name,
                        ..Default::default()
                    }),
                );
                self.input_id += 1;
                let field_value = FieldValue::new(expr);
                memory.push(Data::Bool(BoolData(field_value.into())))
            }
            SupportedType::Curve(r) => {
                let expr = Expr::Curve(CurveExpr::Input(
                    self.input_id,
                    Rc::new(curve_expr::InputInfo {
                        kind: InputKind::Plaintext,
                        name,
                        ..Default::default()
                    }),
                ));
                self.input_id += 1;
                let curve_value = CurveValue::from_expr(expr);
                memory.push(Data::Curve(CurveData(curve_value, r)))
            }
            SupportedType::Struct(s) => {
                let field_types = s.types();
                let v: Result<Vec<_>, _> = field_types
                    .iter()
                    .enumerate()
                    .map(|(field_idx, t)| {
                        let mut name = name.clone();
                        name.push('.');
                        name.push_str(&s.get_field_name(field_idx));
                        self.create_inputs(memory, t.clone(), name)
                    })
                    .collect();
                let v = v?;
                memory.push(Data::Collection(v, CollectionKind::Struct(s.clone())))
            }
            SupportedType::Ref(ty) => {
                let object = self.create_inputs(memory, *ty, name)?;
                memory.push(object.as_data())
            }
            SupportedType::Array(ty, n) => {
                let ArrayLen::Computed(n) = n else {
                    return Err("Array length is unknown.");
                };
                let data_vec: Result<Vec<_>, _> = (0..n)
                    .map(|i| {
                        let mut name = name.clone();
                        name.push_str(&format!("[{}]", i));
                        self.create_inputs(memory, *ty.clone(), name)
                    })
                    .collect();
                let data_vec = data_vec?;
                let data = Data::Collection(data_vec, CollectionKind::Array(*ty.clone()));
                memory.push(data)
            }
            SupportedType::Tuple(v) => {
                let data_vec: Result<Vec<_>, _> = v
                    .into_iter()
                    .enumerate()
                    .map(|(field_idx, x)| {
                        let mut name = name.clone();
                        name.push_str(&format!(".{}", field_idx));
                        self.create_inputs(memory, x, name)
                    })
                    .collect();
                let data_vec = data_vec?;
                let data = Data::Collection(data_vec, CollectionKind::Tuple);
                memory.push(data)
            }
            SupportedType::Slice(ty, len) => {
                let Ok(len) = len.value() else {
                    return Err("Slices cannot be inputs.");
                };
                self.create_inputs(
                    memory,
                    SupportedType::Array(ty, ArrayLen::Computed(len)),
                    name,
                )?
            }
            SupportedType::Fn => return Err("Functions cannot be inputs."),
            SupportedType::ImplAny => {
                return Err("Uninferred types like `impl X` cannot be inputs.")
            }
            SupportedType::Var(_) => {
                return Err("Uninferred types like generic types cannot be inputs.")
            }
        };
        Ok(address)
    }
}
