//! Session management and heartbeat

use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::RwLock;
use uuid::Uuid;

/// Session state
#[derive(Debug, Clone)]
pub struct Session {
    pub user_id: Uuid,
    pub session_id: Uuid,
    pub created_at: Instant,
    pub last_heartbeat: Instant,
    pub bytes_tx: u64,
    pub bytes_rx: u64,
    pub sni: String,
}

impl Session {
    pub fn new(user_id: Uuid, sni: String) -> Self {
        let now = Instant::now();
        Self {
            user_id,
            session_id: Uuid::new_v4(),
            created_at: now,
            last_heartbeat: now,
            bytes_tx: 0,
            bytes_rx: 0,
            sni,
        }
    }

    pub fn heartbeat(&mut self) {
        self.last_heartbeat = Instant::now();
    }

    pub fn is_alive(&self, timeout: Duration) -> bool {
        self.last_heartbeat.elapsed() < timeout
    }
}

/// Session manager
pub struct SessionManager {
    sessions: Arc<RwLock<HashMap<Uuid, Session>>>,
    heartbeat_timeout: Duration,
}

impl SessionManager {
    pub fn new(heartbeat_timeout_secs: u64) -> Self {
        Self {
            sessions: Arc::new(RwLock::new(HashMap::new())),
            heartbeat_timeout: Duration::from_secs(heartbeat_timeout_secs),
        }
    }

    pub async fn create_session(&self, user_id: Uuid, sni: String) -> Uuid {
        let session = Session::new(user_id, sni);
        let session_id = session.session_id;
        self.sessions.write().await.insert(session_id, session);
        session_id
    }

    pub async fn heartbeat(&self, session_id: &Uuid) -> bool {
        if let Some(session) = self.sessions.write().await.get_mut(session_id) {
            session.heartbeat();
            true
        } else {
            false
        }
    }

    pub async fn get_session(&self, session_id: &Uuid) -> Option<Session> {
        self.sessions.read().await.get(session_id).cloned()
    }

    pub async fn remove_session(&self, session_id: &Uuid) -> Option<Session> {
        self.sessions.write().await.remove(session_id)
    }

    pub async fn cleanup_stale(&self) -> Vec<Uuid> {
        let mut sessions = self.sessions.write().await;
        let timeout = self.heartbeat_timeout;
        let stale: Vec<Uuid> = sessions
            .iter()
            .filter(|(_, s)| !s.is_alive(timeout))
            .map(|(id, _)| *id)
            .collect();

        for id in &stale {
            sessions.remove(id);
        }
        stale
    }

    pub async fn update_traffic(&self, session_id: &Uuid, tx: u64, rx: u64) {
        if let Some(session) = self.sessions.write().await.get_mut(session_id) {
            session.bytes_tx += tx;
            session.bytes_rx += rx;
        }
    }
}

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