use std::sync::Arc;

use arrow_array::{ArrayRef, Float64Array, ListArray, StructArray};
use arrow_buffer::OffsetBuffer;
use arrow_schema::{DataType, Field};
use derive_more::From;
use indexmap::IndexMap;
use num::complex::Complex64;

use crate::dataset::schema::{
    arrays::{
        ComplexArray, DatasetArray, FixedStepTraceArray, ScalarArray, ScalarListArray,
        VariableStepTraceArray,
    },
    error::DatasetError,
    model::{DatasetDataType, DatasetSchema, ScalarKind, TraceKind},
};

#[derive(Debug, Clone)]
pub struct FixedStepTrace {
    x0: f64,
    step: f64,
    y: ScalarArray,
}

impl FixedStepTrace {
    #[must_use]
    pub fn new(x0: f64, step: f64, y: ScalarArray) -> Self {
        Self { x0, step, y }
    }
    #[must_use]
    pub fn x0(&self) -> f64 {
        self.x0
    }
    #[must_use]
    pub fn step(&self) -> f64 {
        self.step
    }
    #[must_use]
    pub fn y(&self) -> &ScalarArray {
        &self.y
    }
    #[must_use]
    pub fn scalar_kind(&self) -> ScalarKind {
        self.y.scalar_kind()
    }
}

#[derive(Debug, Clone)]
pub struct VariableStepTrace {
    x: Arc<Float64Array>,
    y: ScalarArray,
}

impl VariableStepTrace {
    pub fn new(x: Arc<Float64Array>, y: ScalarArray) -> Result<Self, DatasetError> {
        if x.len() == y.len() {
            Ok(Self { x, y })
        } else {
            Err(DatasetError::TraceLengthMismatch)
        }
    }
    #[must_use]
    pub fn x(&self) -> &Arc<Float64Array> {
        &self.x
    }
    #[must_use]
    pub fn y(&self) -> &ScalarArray {
        &self.y
    }
    #[must_use]
    pub fn scalar_kind(&self) -> ScalarKind {
        self.y.scalar_kind()
    }
}

#[derive(Debug, Clone, From)]
pub enum DatasetScalar {
    Numeric(f64),
    Complex(Complex64),
    SimpleTrace(ScalarArray),
    FixedStepTrace(FixedStepTrace),
    VariableStepTrace(VariableStepTrace),
}

impl DatasetScalar {
    #[must_use]
    pub fn data_type(&self) -> DatasetDataType {
        match self {
            DatasetScalar::Numeric(_) => DatasetDataType::Scalar(ScalarKind::Numeric),
            DatasetScalar::Complex(_) => DatasetDataType::Scalar(ScalarKind::Complex),
            DatasetScalar::SimpleTrace(t) => {
                DatasetDataType::Trace(TraceKind::Simple, t.scalar_kind())
            }
            DatasetScalar::FixedStepTrace(t) => {
                DatasetDataType::Trace(TraceKind::FixedStep, t.scalar_kind())
            }
            DatasetScalar::VariableStepTrace(t) => {
                DatasetDataType::Trace(TraceKind::VariableStep, t.scalar_kind())
            }
        }
    }
}

pub struct DatasetRow(pub IndexMap<String, DatasetScalar>);

impl DatasetRow {
    #[must_use]
    pub fn to_schema(&self) -> DatasetSchema {
        let columns = self
            .0
            .iter()
            .map(|(name, scalar)| (name.clone(), scalar.data_type()))
            .collect();
        DatasetSchema::new(columns)
    }
}

impl From<FixedStepTrace> for FixedStepTraceArray {
    fn from(value: FixedStepTrace) -> Self {
        let x0: ArrayRef = Arc::new(Float64Array::from(vec![value.x0()]));
        let step: ArrayRef = Arc::new(Float64Array::from(vec![value.step()]));
        let y: ArrayRef = ScalarListArray::from_single_item(value.y().clone()).into();
        let array = StructArray::try_from(vec![("x0", x0), ("step", step), ("y", y)])
            .expect("Should have same length of 1.");
        Self::from_struct_array(Arc::new(array), value.scalar_kind())
    }
}

impl From<VariableStepTrace> for VariableStepTraceArray {
    fn from(value: VariableStepTrace) -> Self {
        let x_field = Arc::new(Field::new_list_field(DataType::Float64, false));
        let x_offsets = OffsetBuffer::from_lengths([value.x().len()]);
        let x: ArrayRef = Arc::new(ListArray::new(x_field, x_offsets, value.x().clone(), None));
        let y: ArrayRef = ScalarListArray::from_single_item(value.y().clone()).into();
        let array =
            StructArray::try_from(vec![("x", x), ("y", y)]).expect("Should have same length of 1.");
        Self::from_struct_array(Arc::new(array), value.scalar_kind())
    }
}

impl From<DatasetScalar> for DatasetArray {
    fn from(value: DatasetScalar) -> Self {
        match value {
            DatasetScalar::Numeric(v) => Arc::new(Float64Array::from(vec![v])).into(),
            DatasetScalar::Complex(v) => ComplexArray::from_iter(vec![v]).into(),
            DatasetScalar::SimpleTrace(v) => ScalarListArray::from_single_item(v).into(),
            DatasetScalar::FixedStepTrace(v) => FixedStepTraceArray::from(v).into(),
            DatasetScalar::VariableStepTrace(v) => VariableStepTraceArray::from(v).into(),
        }
    }
}
