use std::sync::{
    Arc, LazyLock,
    atomic::{AtomicU8, Ordering},
};

use crate::utils::string_generator::NonceGenerator;

// MARK: - Trackable

pub(crate) trait Trackable {
    fn encode_query(&self) -> String;
    fn state_cell(&self) -> &Arc<TrackingStateCell>;
}

// MARK: - TrackingState

#[derive(Debug, Clone, Copy, PartialEq)]
pub(crate) enum TrackingState {
    Unsent = 0,
    Transmitting = 1,
    Sent = 2,
}

impl From<u8> for TrackingState {
    fn from(value: u8) -> Self {
        match value {
            0 => Self::Unsent,
            1 => Self::Transmitting,
            2 => Self::Sent,
            _ => Self::Unsent,
        }
    }
}

// MARK: - TrackingStateCell

#[derive(Debug)]
pub(crate) enum TrackingStateCell {
    Once(AtomicU8),
    Resendable,
}

impl TrackingStateCell {
    pub fn once() -> Self {
        return Self::Once(AtomicU8::new(TrackingState::Unsent as u8));
    }

    pub fn state(&self) -> TrackingState {
        match self {
            Self::Once(state) => state.load(Ordering::Relaxed).into(),
            Self::Resendable => TrackingState::Unsent,
        }
    }

    pub fn change_state(&self, new_state: TrackingState) -> bool {
        match self {
            Self::Once(state) => {
                let new_state = new_state as u8;
                if new_state == TrackingState::Sent as u8 {
                    // Any (0|1|2) -> Sent (2)
                    state.store(new_state, Ordering::Relaxed);
                    true
                } else {
                    // Unsent (0) <=> Transmitting (1)
                    let required_state = 1 - new_state;
                    state.compare_exchange(required_state, new_state, Ordering::Relaxed, Ordering::Relaxed).is_ok()
                }
            }
            Self::Resendable => false,
        }
    }
}

impl Clone for TrackingStateCell {
    fn clone(&self) -> Self {
        match self {
            Self::Once(state) => Self::Once(AtomicU8::new(state.load(Ordering::Relaxed))),
            Self::Resendable => Self::Resendable,
        }
    }
}

#[cfg(test)]
impl PartialEq for TrackingStateCell {
    fn eq(&self, other: &Self) -> bool {
        match (self, other) {
            (Self::Once(a), Self::Once(b)) => a.load(Ordering::Relaxed) == b.load(Ordering::Relaxed),
            (Self::Resendable, Self::Resendable) => true,
            _ => false,
        }
    }
}

// MARK: - TrackableBase

#[derive(Debug)]
pub(crate) struct TrackableBase {
    nonce: LazyLock<String>,
    state: Arc<TrackingStateCell>,
}

impl TrackableBase {
    pub fn new(resendable: bool) -> Self {
        let state = if resendable { TrackingStateCell::Resendable } else { TrackingStateCell::once() };
        Self {
            nonce: LazyLock::new(NonceGenerator::generate),
            state: Arc::new(state),
        }
    }

    pub fn nonce(&self) -> &str {
        return &self.nonce;
    }

    pub fn state_cell(&self) -> &Arc<TrackingStateCell> {
        return &self.state;
    }
}

impl Default for TrackableBase {
    fn default() -> Self {
        return Self::new(false);
    }
}

#[cfg(test)]
impl Clone for TrackableBase {
    fn clone(&self) -> Self {
        Self {
            nonce: LazyLock::new(NonceGenerator::generate),
            state: self.state.clone(),
        }
    }
}

#[cfg(test)]
impl PartialEq for TrackableBase {
    fn eq(&self, other: &Self) -> bool {
        // Ignoring nonce values
        return self.state == other.state;
    }
}

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

    #[rstest]
    #[case(TrackingState::Unsent, TrackingState::Unsent, false)]
    #[case(TrackingState::Unsent, TrackingState::Transmitting, true)]
    #[case(TrackingState::Unsent, TrackingState::Sent, true)]
    #[case(TrackingState::Transmitting, TrackingState::Unsent, true)]
    #[case(TrackingState::Transmitting, TrackingState::Transmitting, false)]
    #[case(TrackingState::Transmitting, TrackingState::Sent, true)]
    #[case(TrackingState::Sent, TrackingState::Unsent, false)]
    #[case(TrackingState::Sent, TrackingState::Transmitting, false)]
    #[case(TrackingState::Sent, TrackingState::Sent, true)]
    fn test_once_state_cell_change_state(
        #[case] initial_state: TrackingState,
        #[case] target_state: TrackingState,
        #[case] expect_change: bool,
    ) {
        let expected_state = if expect_change { target_state } else { initial_state };
        let state_cell = TrackingStateCell::once();
        match &state_cell {
            TrackingStateCell::Once(state) => state.store(initial_state as u8, Ordering::Relaxed),
            TrackingStateCell::Resendable => panic!("Expected 'state_cell' to be 'Once', got 'Resendable'"),
        }
        assert_eq!(expect_change, state_cell.change_state(target_state));
        assert_eq!(expected_state, state_cell.state());
    }

    #[rstest]
    #[case(TrackingState::Unsent)]
    #[case(TrackingState::Transmitting)]
    #[case(TrackingState::Sent)]
    fn test_resendable_state_cell_change_state(#[case] target_state: TrackingState) {
        let state_cell = TrackingStateCell::Resendable;
        assert!(!state_cell.change_state(target_state));
        assert_eq!(TrackingState::Unsent, state_cell.state());
    }
}
