use std::fmt;
use std::sync::atomic::{AtomicU8, Ordering};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ConnectionState {
    Initial = 0,
    Connecting = 1,
    Handshake = 2,
    Established = 3,
    Active = 4,
    Reconnecting = 5,
    Disconnecting = 6,
    Disconnected = 7,
}

impl ConnectionState {
    pub fn from_u8(v: u8) -> Self {
        match v {
            0 => ConnectionState::Initial,
            1 => ConnectionState::Connecting,
            2 => ConnectionState::Handshake,
            3 => ConnectionState::Established,
            4 => ConnectionState::Active,
            5 => ConnectionState::Reconnecting,
            6 => ConnectionState::Disconnecting,
            7 => ConnectionState::Disconnected,
            _ => ConnectionState::Initial,
        }
    }

    pub fn as_u8(self) -> u8 {
        self as u8
    }

    pub fn is_connected(self) -> bool {
        matches!(self, ConnectionState::Established | ConnectionState::Active)
    }

    pub fn can_send(self) -> bool {
        matches!(self, ConnectionState::Established | ConnectionState::Active)
    }

    pub fn can_receive(self) -> bool {
        matches!(self, ConnectionState::Established | ConnectionState::Active)
    }

    pub fn should_heartbeat(self) -> bool {
        matches!(self, ConnectionState::Active)
    }
}

impl fmt::Display for ConnectionState {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            ConnectionState::Initial => write!(f, "Initial"),
            ConnectionState::Connecting => write!(f, "Connecting"),
            ConnectionState::Handshake => write!(f, "Handshake"),
            ConnectionState::Established => write!(f, "Established"),
            ConnectionState::Active => write!(f, "Active"),
            ConnectionState::Reconnecting => write!(f, "Reconnecting"),
            ConnectionState::Disconnecting => write!(f, "Disconnecting"),
            ConnectionState::Disconnected => write!(f, "Disconnected"),
        }
    }
}

pub struct ConnectionStateMachine {
    state: AtomicU8,
}

impl ConnectionStateMachine {
    pub fn new() -> Self {
        Self {
            state: AtomicU8::new(ConnectionState::Initial.as_u8()),
        }
    }

    #[cfg(test)]
    pub fn new_with_state(state: ConnectionState) -> Self {
        Self {
            state: AtomicU8::new(state.as_u8()),
        }
    }

    pub fn current(&self) -> ConnectionState {
        ConnectionState::from_u8(self.state.load(Ordering::SeqCst))
    }

    pub fn set(&self, state: ConnectionState) -> bool {
        let current = self.current();
        if Self::can_transition(current, state) {
            self.state.store(state.as_u8(), Ordering::SeqCst);
            true
        } else {
            false
        }
    }

    pub fn transition(&self, to: ConnectionState) -> bool {
        self.set(to)
    }

    fn can_transition(from: ConnectionState, to: ConnectionState) -> bool {
        match (from, to) {
            (ConnectionState::Initial, ConnectionState::Connecting) => true,
            (ConnectionState::Connecting, ConnectionState::Handshake) => true,
            (ConnectionState::Handshake, ConnectionState::Established) => true,
            (ConnectionState::Handshake, ConnectionState::Disconnecting) => true,
            (ConnectionState::Established, ConnectionState::Active) => true,
            (ConnectionState::Established, ConnectionState::Disconnecting) => true,
            (ConnectionState::Active, ConnectionState::Reconnecting) => true,
            (ConnectionState::Active, ConnectionState::Disconnecting) => true,
            (ConnectionState::Reconnecting, ConnectionState::Connecting) => true,
            (ConnectionState::Reconnecting, ConnectionState::Established) => true,
            (ConnectionState::Reconnecting, ConnectionState::Disconnected) => true,
            (ConnectionState::Disconnecting, ConnectionState::Disconnected) => true,
            (ConnectionState::Disconnected, ConnectionState::Connecting) => true,
            _ => false,
        }
    }

    pub fn is_connected(&self) -> bool {
        self.current().is_connected()
    }

    pub fn is_active(&self) -> bool {
        self.current() == ConnectionState::Active
    }

    pub fn should_heartbeat(&self) -> bool {
        self.current().should_heartbeat()
    }
}

impl Default for ConnectionStateMachine {
    fn default() -> Self {
        Self::new()
    }
}
