use crate::{
    core::global_value::{curve_value::CurveValue, field_array::FieldArray, value::FieldValue},
    traits::{Reveal, ToMontgomery},
    utils::{
        curve_point::Curve,
        field::{BaseField, ScalarField},
    },
};
use std::ops::{Add, Index, Mul, Neg, Sub};

#[derive(Debug, Clone, Copy)]
pub struct CurveArray<const N: usize>([CurveValue; N]);

impl<const N: usize> IntoIterator for CurveArray<N> {
    type Item = CurveValue;
    type IntoIter = <[CurveValue; N] as IntoIterator>::IntoIter;

    fn into_iter(self) -> Self::IntoIter {
        self.0.into_iter()
    }
}

impl<const N: usize> Index<usize> for CurveArray<N> {
    type Output = CurveValue;

    fn index(&self, index: usize) -> &Self::Output {
        &self.0[index]
    }
}

impl<const N: usize> Add for CurveArray<N> {
    type Output = Self;

    fn add(self, rhs: Self) -> Self::Output {
        Self(
            self.into_iter()
                .zip(rhs)
                .map(|(val_lhs, val_rhs)| val_lhs + val_rhs)
                .collect::<Vec<CurveValue>>()
                .try_into()
                .unwrap_or_else(|v: Vec<CurveValue>| {
                    panic!("Expected a Vec of length {} (found {})", N, v.len())
                }),
        )
    }
}

impl<const N: usize> Sub for CurveArray<N> {
    type Output = Self;

    fn sub(self, rhs: Self) -> Self::Output {
        Self(
            self.into_iter()
                .zip(rhs)
                .map(|(val_lhs, val_rhs)| val_lhs - val_rhs)
                .collect::<Vec<CurveValue>>()
                .try_into()
                .unwrap_or_else(|v: Vec<CurveValue>| {
                    panic!("Expected a Vec of length {} (found {})", N, v.len())
                }),
        )
    }
}

impl<const N: usize> Mul<CurveArray<N>> for FieldArray<N, ScalarField> {
    type Output = CurveArray<N>;

    fn mul(self, rhs: CurveArray<N>) -> Self::Output {
        CurveArray(
            self.into_iter()
                .zip(rhs)
                .map(|(val_lhs, val_rhs)| val_lhs * val_rhs)
                .collect::<Vec<CurveValue>>()
                .try_into()
                .unwrap_or_else(|v: Vec<CurveValue>| {
                    panic!("Expected a Vec of length {} (found {})", N, v.len())
                }),
        )
    }
}

impl<const N: usize> Neg for CurveArray<N> {
    type Output = Self;

    fn neg(self) -> Self::Output {
        Self(self.0.map(|val| -val))
    }
}

impl<const N: usize> Reveal for CurveArray<N> {
    fn reveal(self) -> Self {
        Self(self.0.map(|val| val.reveal()))
    }
}

impl<const N: usize> ToMontgomery for CurveArray<N> {
    type Output = FieldArray<N, BaseField>;

    fn to_montgomery(
        self,
        is_expected_non_identity: bool,
    ) -> (FieldArray<N, BaseField>, FieldArray<N, BaseField>) {
        let res: (Vec<FieldValue<BaseField>>, Vec<FieldValue<BaseField>>) = self
            .0
            .map(|val| val.to_montgomery(is_expected_non_identity))
            .into_iter()
            .unzip();
        (
            FieldArray::from(
                TryInto::<[FieldValue<BaseField>; N]>::try_into(res.0).unwrap_or_else(
                    |v: Vec<FieldValue<BaseField>>| {
                        panic!("Expected a Vec of length {} (found {})", N, v.len())
                    },
                ),
            ),
            FieldArray::from(
                TryInto::<[FieldValue<BaseField>; N]>::try_into(res.1).unwrap_or_else(
                    |v: Vec<FieldValue<BaseField>>| {
                        panic!("Expected a Vec of length {} (found {})", N, v.len())
                    },
                ),
            ),
        )
    }
}

impl<const N: usize> From<[CurveValue; N]> for CurveArray<N> {
    fn from(value: [CurveValue; N]) -> Self {
        Self(value)
    }
}

impl<const N: usize> Curve for CurveArray<N> {
    fn generator() -> Self {
        Self([CurveValue::generator(); N])
    }
}
