use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};

use async_trait::async_trait;
use dashmap::DashMap;

use crate::socket::session_store::{SessionError, SessionInfo, SessionStoreTrait};

pub struct InMemorySessionStore {
    sessions: Arc<DashMap<String, SessionInfo>>,
}

impl InMemorySessionStore {
    pub fn new() -> Self {
        Self {
            sessions: Arc::new(DashMap::new()),
        }
    }
}

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

fn now_millis() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_millis() as u64
}

#[async_trait]
impl SessionStoreTrait for InMemorySessionStore {
    async fn create(
        &self,
        sid: &str,
        transport: &str,
        server_id: &str,
    ) -> Result<(), SessionError> {
        let info = SessionInfo {
            sid: sid.to_string(),
            transport: transport.to_string(),
            state: "connecting".to_string(),
            server_id: server_id.to_string(),
            created_at: now_millis(),
            last_ping: now_millis(),
        };
        self.sessions.insert(sid.to_string(), info);
        Ok(())
    }

    async fn get(&self, sid: &str) -> Result<Option<SessionInfo>, SessionError> {
        Ok(self.sessions.get(sid).map(|r| r.value().clone()))
    }

    async fn set_state(&self, sid: &str, state: &str) -> Result<(), SessionError> {
        if let Some(mut entry) = self.sessions.get_mut(sid) {
            entry.value_mut().state = state.to_string();
            Ok(())
        } else {
            Err(SessionError::NotFound(sid.to_string()))
        }
    }

    async fn set_transport(&self, sid: &str, transport: &str) -> Result<(), SessionError> {
        if let Some(mut entry) = self.sessions.get_mut(sid) {
            entry.value_mut().transport = transport.to_string();
            Ok(())
        } else {
            Err(SessionError::NotFound(sid.to_string()))
        }
    }

    async fn update_ping(&self, sid: &str) -> Result<(), SessionError> {
        if let Some(mut entry) = self.sessions.get_mut(sid) {
            entry.value_mut().last_ping = now_millis();
            Ok(())
        } else {
            Err(SessionError::NotFound(sid.to_string()))
        }
    }

    async fn remove(&self, sid: &str) -> Result<(), SessionError> {
        self.sessions.remove(sid);
        Ok(())
    }

    async fn exists(&self, sid: &str) -> Result<bool, SessionError> {
        Ok(self.sessions.contains_key(sid))
    }
}