use pyo3::exceptions::asyncio::InvalidStateError;
use pyo3::prelude::*;
use pyo3::types::PyFunction;

use crate::{
    communication::{append_usize, append_usize_vec, retrieve_usize},
    PyAnySerde,
};

#[derive(Clone)]
pub struct UnionSerde {
    pub option_serdes: Vec<Box<dyn PyAnySerde>>,
    pub option_choice_fn: Py<PyFunction>,
}

impl PyAnySerde for UnionSerde {
    fn append<'py>(
        &mut self,
        buf: &mut [u8],
        offset: usize,
        obj: &Bound<'py, PyAny>,
    ) -> PyResult<usize> {
        let serde_idx = self
            .option_choice_fn
            .bind(obj.py())
            .call1((obj,))?
            .extract::<usize>()?;
        let offset = append_usize(buf, offset, serde_idx);
        let pyany_serde = self.option_serdes.get_mut(serde_idx).ok_or_else(|| {
            InvalidStateError::new_err(format!(
                "Serde choice function returned {serde_idx} which is not a valid choice index"
            ))
        })?;
        pyany_serde.append(buf, offset, obj)
    }

    fn append_vec<'py>(
        &mut self,
        v: &mut Vec<u8>,
        start_addr: Option<usize>,
        obj: &Bound<'py, PyAny>,
    ) -> PyResult<()> {
        let serde_idx = self
            .option_choice_fn
            .bind(obj.py())
            .call1((obj,))?
            .extract::<usize>()?;
        append_usize_vec(v, serde_idx);
        let pyany_serde = self.option_serdes.get_mut(serde_idx).ok_or_else(|| {
            InvalidStateError::new_err(format!(
                "Serde choice function returned {serde_idx} which is not a valid choice index"
            ))
        })?;
        pyany_serde.append_vec(v, start_addr, obj)
    }

    fn retrieve<'py>(
        &mut self,
        py: Python<'py>,
        buf: &[u8],
        offset: usize,
    ) -> PyResult<(Bound<'py, PyAny>, usize)> {
        let (serde_idx, offset) = retrieve_usize(buf, offset)?;
        let pyany_serde = self.option_serdes.get_mut(serde_idx).ok_or_else(|| {
            InvalidStateError::new_err(format!(
                "Deserialized serde idx {serde_idx} which is not a valid choice index"
            ))
        })?;
        pyany_serde.retrieve(py, buf, offset)
    }
}
