use std::sync::{
    atomic::{AtomicU32, Ordering},
    Arc,
};

use dashmap::DashMap;

use crate::config::IpString;

use super::{KcpConnMgr, KcpSession};

pub struct KcpServerMgr<Session>
where
    Session: KcpSession,
{
    conn_mgr: Arc<KcpConnMgr<Session>>,
    session_id_counter: AtomicU32,
    session_map: DashMap<u32, Arc<Session>>,
}

impl<Session> KcpServerMgr<Session>
where
    Session: KcpSession + 'static,
{
    pub async fn new(
        outer_ip: IpString,
        listen_port: u16,
        max_conn_num: usize,
        udp_mtu: usize,
        recv_buf_len: usize,
    ) -> Result<Self, std::io::Error> {
        let addr = format!("{}:{}", outer_ip.0, listen_port);
        let conn_mgr = KcpConnMgr::new(&addr, max_conn_num, udp_mtu, recv_buf_len).await?;

        Ok(Self {
            conn_mgr: Arc::new(conn_mgr),
            session_id_counter: AtomicU32::new(0),
            session_map: DashMap::new(),
        })
    }

    pub async fn run(self: Arc<Self>) {
        self.conn_mgr.set_server_mgr(self.clone());
        if let Err(err) = self.conn_mgr.clone().run().await {
            tracing::error!("KcpConnMgr run fail {err:?}")
        }
    }

    pub fn add_session(&self, session_id: u32, session: Session) {
        self.session_map.insert(session_id, Arc::new(session));
    }

    pub fn find_session(&self, session_id: u32) -> Option<Arc<Session>> {
        self.session_map.get(&session_id).map(|s| s.clone())
    }

    pub async fn disconnect(&self, session_id: u32, data: u32) {
        let Some(session) = self.find_session(session_id) else {
            tracing::warn!("disconnect find_session failed, session_id={session_id}");
            return;
        };

        session.on_disconnect(data);
        self.conn_mgr.disconnect(session_id, data).await;
    }

    pub fn session_count(&self) -> usize {
        self.session_map.len()
    }

    pub fn gen_new_session_id(&self) -> u32 {
        self.session_id_counter.fetch_add(1, Ordering::SeqCst) + 1
    }
}
