use pyo3::prelude::*;
use pyo3::types::{PyDict, PyList, PySet, PyTuple};

use crate::communication::append_usize_vec;
use crate::create_numpy_pyany_serde;
use crate::{
    common::{
        detect_python_type, get_python_type_byte, retrieve_python_type, NumpyDtype, PythonType,
    },
    communication::{append_usize, retrieve_usize},
    PyAnySerde,
};

use super::{
    BoolSerde, BytesSerde, ComplexSerde, FloatSerde, IntSerde, NumpySerde, NumpySerdeConfig,
    PickleSerde, StringSerde,
};

#[derive(Clone)]
pub struct DynamicSerde {
    pickle_serde: PickleSerde,
    int_serde: IntSerde,
    float_serde: FloatSerde,
    complex_serde: ComplexSerde,
    boolean_serde: BoolSerde,
    string_serde: StringSerde,
    bytes_serde: BytesSerde,
    numpy_i8_serde: NumpySerde<i8>,
    numpy_i16_serde: NumpySerde<i16>,
    numpy_i32_serde: NumpySerde<i32>,
    numpy_i64_serde: NumpySerde<i64>,
    numpy_u8_serde: NumpySerde<u8>,
    numpy_u16_serde: NumpySerde<u16>,
    numpy_u32_serde: NumpySerde<u32>,
    numpy_u64_serde: NumpySerde<u64>,
    numpy_f32_serde: NumpySerde<f32>,
    numpy_f64_serde: NumpySerde<f64>,
}

impl DynamicSerde {
    pub fn new() -> PyResult<Self> {
        let pickle_serde = PickleSerde::new()?;
        let int_serde = IntSerde {};
        let float_serde = FloatSerde {};
        let complex_serde = ComplexSerde {};
        let boolean_serde = BoolSerde {};
        let string_serde = StringSerde {};
        let bytes_serde = BytesSerde {};
        let numpy_serde_config = NumpySerdeConfig::DYNAMIC {
            preprocessor_fn: None,
            postprocessor_fn: None,
        };
        let numpy_i8_serde = *create_numpy_pyany_serde!(i8, numpy_serde_config.clone());
        let numpy_i16_serde = *create_numpy_pyany_serde!(i16, numpy_serde_config.clone());
        let numpy_i32_serde = *create_numpy_pyany_serde!(i32, numpy_serde_config.clone());
        let numpy_i64_serde = *create_numpy_pyany_serde!(i64, numpy_serde_config.clone());
        let numpy_u8_serde = *create_numpy_pyany_serde!(u8, numpy_serde_config.clone());
        let numpy_u16_serde = *create_numpy_pyany_serde!(u16, numpy_serde_config.clone());
        let numpy_u32_serde = *create_numpy_pyany_serde!(u32, numpy_serde_config.clone());
        let numpy_u64_serde = *create_numpy_pyany_serde!(u64, numpy_serde_config.clone());
        let numpy_f32_serde = *create_numpy_pyany_serde!(f32, numpy_serde_config.clone());
        let numpy_f64_serde = *create_numpy_pyany_serde!(f64, numpy_serde_config.clone());

        Ok(DynamicSerde {
            pickle_serde,
            int_serde,
            float_serde,
            complex_serde,
            boolean_serde,
            string_serde,
            bytes_serde,
            numpy_i8_serde,
            numpy_i16_serde,
            numpy_i32_serde,
            numpy_i64_serde,
            numpy_u8_serde,
            numpy_u16_serde,
            numpy_u32_serde,
            numpy_u64_serde,
            numpy_f32_serde,
            numpy_f64_serde,
        })
    }
}

impl PyAnySerde for DynamicSerde {
    fn append<'py>(
        &mut self,
        buf: &mut [u8],
        mut offset: usize,
        obj: &Bound<'py, PyAny>,
    ) -> PyResult<usize> {
        let python_type = detect_python_type(obj)?;
        buf[offset] = get_python_type_byte(&python_type);
        match python_type {
            PythonType::BOOL => {
                offset = self.boolean_serde.append(buf, offset, obj)?;
            }
            PythonType::INT => {
                offset = self.int_serde.append(buf, offset, obj)?;
            }
            PythonType::FLOAT => {
                offset = self.float_serde.append(buf, offset, obj)?;
            }
            PythonType::COMPLEX => {
                offset = self.complex_serde.append(buf, offset, obj)?;
            }
            PythonType::STRING => {
                offset = self.string_serde.append(buf, offset, obj)?;
            }
            PythonType::BYTES => {
                offset = self.bytes_serde.append(buf, offset, obj)?;
            }
            PythonType::NUMPY { dtype } => match dtype {
                NumpyDtype::INT8 => {
                    offset = self.numpy_i8_serde.append(buf, offset, obj)?;
                }
                NumpyDtype::INT16 => {
                    offset = self.numpy_i16_serde.append(buf, offset, obj)?;
                }
                NumpyDtype::INT32 => {
                    offset = self.numpy_i32_serde.append(buf, offset, obj)?;
                }
                NumpyDtype::INT64 => {
                    offset = self.numpy_i64_serde.append(buf, offset, obj)?;
                }
                NumpyDtype::UINT8 => {
                    offset = self.numpy_u8_serde.append(buf, offset, obj)?;
                }
                NumpyDtype::UINT16 => {
                    offset = self.numpy_u16_serde.append(buf, offset, obj)?;
                }
                NumpyDtype::UINT32 => {
                    offset = self.numpy_u32_serde.append(buf, offset, obj)?;
                }
                NumpyDtype::UINT64 => {
                    offset = self.numpy_u64_serde.append(buf, offset, obj)?;
                }
                NumpyDtype::FLOAT32 => {
                    offset = self.numpy_f32_serde.append(buf, offset, obj)?;
                }
                NumpyDtype::FLOAT64 => {
                    offset = self.numpy_f64_serde.append(buf, offset, obj)?;
                }
            },
            PythonType::LIST => {
                let list = obj.cast::<PyList>()?;
                offset = append_usize(buf, offset, list.len());
                for item in list.iter() {
                    offset = self.append(buf, offset, &item)?;
                }
            }
            PythonType::SET => {
                let set = obj.cast::<PySet>()?;
                offset = append_usize(buf, offset, set.len());
                for item in set.iter() {
                    offset = self.append(buf, offset, &item)?;
                }
            }
            PythonType::TUPLE => {
                let tuple = obj.cast::<PyTuple>()?;
                offset = append_usize(buf, offset, tuple.len());
                for item in tuple.iter() {
                    offset = self.append(buf, offset, &item)?;
                }
            }
            PythonType::DICT => {
                let dict = obj.cast::<PyDict>()?;
                offset = append_usize(buf, offset, dict.len());
                for (key, value) in dict.iter() {
                    offset = self.append(buf, offset, &key)?;
                    offset = self.append(buf, offset, &value)?;
                }
            }
            PythonType::OTHER => {
                offset = self.pickle_serde.append(buf, offset, obj)?;
            }
        };
        Ok(offset)
    }

    fn append_vec<'py>(
        &mut self,
        v: &mut Vec<u8>,
        start_addr: Option<usize>,
        obj: &Bound<'py, PyAny>,
    ) -> PyResult<()> {
        let python_type = detect_python_type(obj)?;
        v.push(get_python_type_byte(&python_type));
        match python_type {
            PythonType::BOOL => {
                self.boolean_serde.append_vec(v, start_addr, obj)?;
            }
            PythonType::INT => {
                self.int_serde.append_vec(v, start_addr, obj)?;
            }
            PythonType::FLOAT => {
                self.float_serde.append_vec(v, start_addr, obj)?;
            }
            PythonType::COMPLEX => {
                self.complex_serde.append_vec(v, start_addr, obj)?;
            }
            PythonType::STRING => {
                self.string_serde.append_vec(v, start_addr, obj)?;
            }
            PythonType::BYTES => {
                self.bytes_serde.append_vec(v, start_addr, obj)?;
            }
            PythonType::NUMPY { dtype } => match dtype {
                NumpyDtype::INT8 => {
                    self.numpy_i8_serde.append_vec(v, start_addr, obj)?;
                }
                NumpyDtype::INT16 => {
                    self.numpy_i16_serde.append_vec(v, start_addr, obj)?;
                }
                NumpyDtype::INT32 => {
                    self.numpy_i32_serde.append_vec(v, start_addr, obj)?;
                }
                NumpyDtype::INT64 => {
                    self.numpy_i64_serde.append_vec(v, start_addr, obj)?;
                }
                NumpyDtype::UINT8 => {
                    self.numpy_u8_serde.append_vec(v, start_addr, obj)?;
                }
                NumpyDtype::UINT16 => {
                    self.numpy_u16_serde.append_vec(v, start_addr, obj)?;
                }
                NumpyDtype::UINT32 => {
                    self.numpy_u32_serde.append_vec(v, start_addr, obj)?;
                }
                NumpyDtype::UINT64 => {
                    self.numpy_u64_serde.append_vec(v, start_addr, obj)?;
                }
                NumpyDtype::FLOAT32 => {
                    self.numpy_f32_serde.append_vec(v, start_addr, obj)?;
                }
                NumpyDtype::FLOAT64 => {
                    self.numpy_f64_serde.append_vec(v, start_addr, obj)?;
                }
            },
            PythonType::LIST => {
                let list = obj.cast::<PyList>()?;
                append_usize_vec(v, list.len());
                for item in list.iter() {
                    self.append_vec(v, start_addr, &item)?;
                }
            }
            PythonType::SET => {
                let set = obj.cast::<PyList>()?;
                append_usize_vec(v, set.len());
                for item in set.iter() {
                    self.append_vec(v, start_addr, &item)?;
                }
            }
            PythonType::TUPLE => {
                let tuple = obj.cast::<PyList>()?;
                append_usize_vec(v, tuple.len());
                for item in tuple.iter() {
                    self.append_vec(v, start_addr, &item)?;
                }
            }
            PythonType::DICT => {
                let dict = obj.cast::<PyDict>()?;
                append_usize_vec(v, dict.len());
                for (key, value) in dict.iter() {
                    self.append_vec(v, start_addr, &key)?;
                    self.append_vec(v, start_addr, &value)?;
                }
            }
            PythonType::OTHER => {
                self.pickle_serde.append_vec(v, start_addr, obj)?;
            }
        };
        Ok(())
    }

    fn retrieve<'py>(
        &mut self,
        py: Python<'py>,
        buf: &[u8],
        offset: usize,
    ) -> PyResult<(Bound<'py, PyAny>, usize)> {
        let (python_type, mut offset) = retrieve_python_type(buf, offset)?;
        let obj;
        match python_type {
            PythonType::BOOL => {
                (obj, offset) = self.boolean_serde.retrieve(py, buf, offset)?;
            }
            PythonType::INT => {
                (obj, offset) = self.int_serde.retrieve(py, buf, offset)?;
            }
            PythonType::FLOAT => {
                (obj, offset) = self.float_serde.retrieve(py, buf, offset)?;
            }
            PythonType::COMPLEX => {
                (obj, offset) = self.complex_serde.retrieve(py, buf, offset)?;
            }
            PythonType::STRING => {
                (obj, offset) = self.string_serde.retrieve(py, buf, offset)?;
            }
            PythonType::BYTES => {
                (obj, offset) = self.bytes_serde.retrieve(py, buf, offset)?;
            }
            PythonType::NUMPY { dtype } => match dtype {
                NumpyDtype::INT8 => {
                    (obj, offset) = self.numpy_i8_serde.retrieve(py, buf, offset)?;
                }
                NumpyDtype::INT16 => {
                    (obj, offset) = self.numpy_i16_serde.retrieve(py, buf, offset)?;
                }
                NumpyDtype::INT32 => {
                    (obj, offset) = self.numpy_i32_serde.retrieve(py, buf, offset)?;
                }
                NumpyDtype::INT64 => {
                    (obj, offset) = self.numpy_i64_serde.retrieve(py, buf, offset)?;
                }
                NumpyDtype::UINT8 => {
                    (obj, offset) = self.numpy_u8_serde.retrieve(py, buf, offset)?;
                }
                NumpyDtype::UINT16 => {
                    (obj, offset) = self.numpy_u16_serde.retrieve(py, buf, offset)?;
                }
                NumpyDtype::UINT32 => {
                    (obj, offset) = self.numpy_u32_serde.retrieve(py, buf, offset)?;
                }
                NumpyDtype::UINT64 => {
                    (obj, offset) = self.numpy_u64_serde.retrieve(py, buf, offset)?;
                }
                NumpyDtype::FLOAT32 => {
                    (obj, offset) = self.numpy_f32_serde.retrieve(py, buf, offset)?;
                }
                NumpyDtype::FLOAT64 => {
                    (obj, offset) = self.numpy_f64_serde.retrieve(py, buf, offset)?;
                }
            },
            PythonType::LIST => {
                let list = PyList::empty(py);
                let n_items;
                (n_items, offset) = retrieve_usize(buf, offset)?;
                for _ in 0..n_items {
                    let item;
                    (item, offset) = self.retrieve(py, buf, offset)?;
                    list.append(item)?;
                }
                obj = list.into_any();
            }
            PythonType::SET => {
                let set = PySet::empty(py)?;
                let n_items;
                (n_items, offset) = retrieve_usize(buf, offset)?;
                for _ in 0..n_items {
                    let item;
                    (item, offset) = self.retrieve(py, buf, offset)?;
                    set.add(item)?;
                }
                obj = set.into_any();
            }
            PythonType::TUPLE => {
                let n_items;
                (n_items, offset) = retrieve_usize(buf, offset)?;
                let mut tuple_vec = Vec::with_capacity(n_items);
                for _ in 0..n_items {
                    let item;
                    (item, offset) = self.retrieve(py, buf, offset)?;
                    tuple_vec.push(item);
                }
                obj = PyTuple::new(py, tuple_vec)?.into_any();
            }
            PythonType::DICT => {
                let dict = PyDict::new(py);
                let n_items;
                (n_items, offset) = retrieve_usize(buf, offset)?;
                for _ in 0..n_items {
                    let key;
                    (key, offset) = self.retrieve(py, buf, offset)?;
                    let value;
                    (value, offset) = self.retrieve(py, buf, offset)?;
                    dict.set_item(key, value)?;
                }
                obj = dict.into_any();
            }
            PythonType::OTHER => {
                (obj, offset) = self.pickle_serde.retrieve(py, buf, offset)?;
            }
        };
        Ok((obj, offset))
    }
}
