use std::{
    error::Error,
    net::SocketAddr,
    sync::{Arc, OnceLock, Weak},
    task::{Context, Poll},
};

use rand::RngCore;
use tokio::net::UdpSocket;

use crate::net::tools::connect_packet::*;

use super::{KcpServerError, KcpServerMgr, KcpSession};

pub struct KcpConnMgr<Session>
where
    Session: KcpSession,
{
    socket: UdpSocket,
    server_mgr: OnceLock<Weak<KcpServerMgr<Session>>>,
    #[allow(dead_code)]
    max_conn_num: usize, // todo
    udp_mtu: usize,
    recv_buf_len: usize,
}

impl<Session> KcpConnMgr<Session>
where
    Session: KcpSession + 'static,
{
    pub async fn new(
        addr: &str,
        max_conn_num: usize,
        udp_mtu: usize,
        recv_buf_len: usize,
    ) -> Result<Self, std::io::Error> {
        let socket = UdpSocket::bind(addr).await?;

        Ok(Self {
            socket,
            server_mgr: OnceLock::new(),
            max_conn_num,
            udp_mtu,
            recv_buf_len,
        })
    }

    pub async fn run(self: Arc<Self>) -> Result<(), Box<dyn Error>> {
        let mut buf = vec![0u8; self.udp_mtu];
        loop {
            let Ok((len, addr)) = self.socket.recv_from(&mut buf).await else {
                continue;
            };

            let buf = &buf[..len];
            match len {
                CONTROL_PACKET_SIZE => match get_control_packet_type(buf) {
                    Some(ControlPacketType::Connect) => {
                        if let Err(err) = self.handle_connect_packet(buf, addr).await {
                            tracing::warn!("handle_connect_packet failed {err:?}");
                        }
                    }
                    Some(ControlPacketType::Disconnect) => {
                        if let Err(err) = self.handle_disconnect_packet(buf, addr).await {
                            tracing::warn!("handle_disconnect_packet failed {err:?}");
                        }
                    }
                    Some(ControlPacketType::SendBackConv) => {
                        tracing::warn!("server received SendBackConv packet from {addr:?}")
                    }
                    None => (),
                },
                28.. => {
                    if let Err(err) = self.handle_kcp_packet(buf.into()).await {
                        tracing::warn!("handle_kcp_packet failed {err:?}");
                    }
                }
                _ => (),
            };
        }
    }

    async fn handle_connect_packet(
        self: &Arc<Self>,
        buf: &[u8],
        addr: SocketAddr,
    ) -> Result<(), std::io::Error> {
        let server_mgr = self.server_mgr();

        let token = match get_token_from_control_packet(buf) {
            0 => rand::thread_rng().next_u32(),
            token => token,
        };
        let session_id = server_mgr.gen_new_session_id();
        let data = get_data_from_control_packet(buf);

        let mut session = Session::new(self.clone(), session_id, token, self.recv_buf_len, addr);
        session.on_connect(data);
        server_mgr.add_session(session_id, session);

        tracing::info!(
            "add_session from ip={} port={} data={} token={}",
            addr.ip(),
            addr.port(),
            data,
            token
        );

        let packet = make_send_back_conv_packet(session_id, token, data);
        self.socket.send_to(&packet, addr).await?;

        Ok(())
    }

    async fn handle_disconnect_packet(
        self: &Arc<Self>,
        buf: &[u8],
        _addr: SocketAddr,
    ) -> Result<(), std::io::Error> {
        let server_mgr = self.server_mgr();

        let conv = get_conv_from_control_packet(buf);
        let data = get_data_from_control_packet(buf);

        server_mgr.disconnect(conv, data).await;
        tracing::info!(
            "disconnect session from client, session_id={} data={}",
            conv,
            data
        );

        Ok(())
    }

    async fn handle_kcp_packet(&self, buf: Box<[u8]>) -> Result<(), Box<dyn Error>> {
        let conv = mhy_kcp::get_conv(&buf);
        let Some(session) = self.server_mgr().find_session(conv) else {
            return Err(KcpServerError::SessionNotFound(conv))?;
        };

        tokio::spawn(async move { session.input(buf).await });
        Ok(())
    }

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

        let packet = make_disconnect_packet(session_id, session.get_user_token(), data);
        let _ = self
            .socket
            .send_to(&packet, session.get_udp_remote_endpoint())
            .await;
    }

    pub fn send_udp_packet(
        &self,
        cx: &mut Context<'_>,
        buf: &[u8],
        endpoint: SocketAddr,
    ) -> Poll<Result<usize, std::io::Error>> {
        self.socket.poll_send_to(cx, buf, endpoint)
    }

    pub fn set_server_mgr(&self, server_mgr: Arc<KcpServerMgr<Session>>) {
        if self.server_mgr.set(Arc::downgrade(&server_mgr)).is_err() {
            panic!("KcpConnMgr::set_server_mgr shouldn't be called twice")
        }
    }

    fn server_mgr(&self) -> Arc<KcpServerMgr<Session>> {
        self.server_mgr.get().unwrap().upgrade().unwrap()
    }
}
