use std::sync::Arc;
use std::time::Duration;

use crate::engine::packet::Packet;
use crate::engine::session::{SessionState, SessionStore, TransportType};

pub struct HeartbeatManager {
    store: SessionStore,
    ping_interval: u64,
    ping_timeout: u64,
}

impl HeartbeatManager {
    pub fn new(store: SessionStore, ping_interval: u64, ping_timeout: u64) -> Self {
        Self {
            store,
            ping_interval,
            ping_timeout,
        }
    }

    pub fn start(self: Arc<Self>) -> tokio::task::JoinHandle<()> {
        let this = self.clone();
        tokio::spawn(async move {
            this.run().await;
        })
    }

    async fn run(&self) {
        let mut interval = tokio::time::interval(Duration::from_millis(self.ping_interval));

        loop {
            interval.tick().await;
            self.check_sessions().await;
        }
    }

    async fn check_sessions(&self) {
        let now = std::time::Instant::now();
        let timeout_duration = Duration::from_millis(self.ping_interval + self.ping_timeout);

        let mut to_remove = Vec::new();

        for entry in self.store.sessions.iter() {
            let sid = entry.key().clone();
            let session = entry.value().clone();

            let (state, last_ping, transport) = {
                let s = session.read().await;
                (s.state, s.last_ping, s.transport)
            };

            if state == SessionState::Closed {
                to_remove.push(sid);
                continue;
            }

            if now.duration_since(last_ping) > timeout_duration {
                tracing::warn!("Session {} timed out", sid);
                to_remove.push(sid);
                continue;
            }

            // For polling sessions: buffer a ping packet for the next GET request.
            // WS/WT sessions rely on their own dedicated ping tasks; the timeout
            // check above already serves as the safety net for all transports.
            if state == SessionState::Open && transport == TransportType::Polling {
                let mut s = session.write().await;
                s.buffer_packet(Packet::ping(""));
            }
        }

        for sid in to_remove {
            self.store.remove(&sid);
        }
    }
}