use log::trace;
use std::any::{Any, type_name};
use std::collections::HashMap;
use std::fmt::{Debug, Formatter};
use std::marker::PhantomData;
use std::sync::RwLock;
use std::sync::atomic::{AtomicUsize, Ordering};

#[derive(Debug)]
pub struct StateVector {
    next_section: AtomicUsize,
    sections: RwLock<HashMap<SectionIdentifier, Box<RwLock<dyn Any + Send + Sync>>>>,
}

#[derive(Clone, Copy, Eq, PartialEq, Hash, Debug)]
pub struct SectionIdentifier(usize);

pub struct SectionWriter<'a, T> {
    id: SectionIdentifier,
    state_vector: &'a StateVector,
    _phantom_data: PhantomData<T>,
}

impl<T> Debug for SectionWriter<'_, T> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "SectionWriter<T={}> {{ id: {:?}, state_vector: {:?} }}",
            type_name::<T>(),
            self.id,
            self.state_vector
        )
    }
}

impl<T: 'static> SectionWriter<'_, T> {
    pub fn get_identifier(&self) -> SectionIdentifier {
        trace!(
            "SectionWriter<T={}>::get_identifier(self: {self:?})",
            type_name::<T>()
        );
        self.id
    }

    pub fn update<F, R>(&self, f: F) -> R
    where
        F: FnOnce(&mut T) -> R,
    {
        trace!(
            "SectionWriter<T={}>::update(self: {self:?}, f)",
            type_name::<T>()
        );
        self.state_vector.sections.clear_poison();
        let sections = self.state_vector.sections.read().unwrap();
        let section = sections.get(&self.id).unwrap();
        let mut data = section.write().unwrap();
        let result = data.downcast_mut::<T>().unwrap();
        f(result)
    }
}

impl StateVector {
    pub fn new() -> Self {
        trace!("StateVector::new()");
        Self {
            next_section: AtomicUsize::new(0usize),
            sections: RwLock::new(HashMap::new()),
        }
    }

    pub fn create_section<T>(&self, initial_value: T) -> SectionWriter<'_, T>
    where
        T: Send + Sync + 'static,
    {
        trace!(
            "StateVector::create_section<T={}>(self: {self:?}, initial_value)",
            type_name::<T>()
        );
        let id = SectionIdentifier(self.next_section.fetch_add(1usize, Ordering::SeqCst));
        let lock = Box::new(RwLock::new(initial_value));

        self.sections.clear_poison();
        let mut sections = self.sections.write().unwrap();
        sections.entry(id).or_insert(lock);

        drop(sections);

        SectionWriter {
            id,
            state_vector: self,
            _phantom_data: PhantomData,
        }
    }

    pub fn access_section<T, F, R>(&self, id: SectionIdentifier, f: F) -> Option<R>
    where
        T: 'static,
        F: FnOnce(&T) -> R,
    {
        trace!(
            "StateVector::access_section<T={}, F={}, R={}>(self: {self:?}, id: {id:?}, f)",
            type_name::<T>(),
            type_name::<F>(),
            type_name::<R>()
        );
        self.sections.clear_poison();
        let Ok(sections) = self.sections.read() else {
            return None;
        };
        let section = sections.get(&id)?;
        section.clear_poison();
        let Ok(data) = section.read() else {
            return None;
        };
        let inner = data.downcast_ref::<T>()?;
        Some(f(inner))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use anyhow::Result;

    #[derive(Default)]
    struct TestType {
        value1: i32,
        value2: i32,
    }

    #[test]
    fn test_two_sections() -> Result<()> {
        let state_vector = StateVector::new();

        let section_1 = state_vector.create_section(TestType::default());
        let section_2 = state_vector.create_section(TestType::default());

        section_1.update(|s| {
            s.value1 = 1;
            s.value2 = 2;
        });
        section_2.update(|s| {
            s.value1 = 3;
            s.value2 = 4;
        });

        let id_1 = section_1.get_identifier();

        state_vector.access_section(id_1, |s: &TestType| {
            assert_eq!(1, s.value1);
            assert_eq!(2, s.value2);
        });

        let id_2 = section_2.get_identifier();

        state_vector.access_section(id_2, |s: &TestType| {
            assert_eq!(3, s.value1);
            assert_eq!(4, s.value2);
        });

        Ok(())
    }
}
