use std::{
    cell::LazyCell,
    collections::BTreeSet,
    net::{IpAddr, SocketAddr},
};

use socket2::{Domain, Protocol, SockAddr, Socket, Type};
use tokio::{
    io::Result as IoResult,
    net::{TcpListener, UdpSocket},
};

use crate::{
    error::Error,
    utils::{ephemeral_port_range, Network},
};

use super::{Connection, udp_sessions::UdpSessions};

const EPHEMERAL_PORTS: LazyCell<BTreeSet<u16>> = LazyCell::new(|| {
    ephemeral_port_range().collect()
});

impl Connection {
    async fn bind_tcp(&self, bind_ip: IpAddr, port: u16, only_v6: Option<bool>) -> IoResult<TcpListener> {
        let domain = match bind_ip {
            IpAddr::V4(_) => Domain::IPV4,
            IpAddr::V6(_) => Domain::IPV6,
        };

        let socket = Socket::new(domain, Type::STREAM, Some(Protocol::TCP))?;

        if let Some(only_v6) = only_v6 {
            socket.set_only_v6(only_v6)?;
        }

        socket.set_nonblocking(true)?;

        socket.bind(&SockAddr::from(
            SocketAddr::new(bind_ip, port)
        ))?;

        // backlog is the queue length for pending connections
        socket.listen(128)?;

        TcpListener::from_std(socket.into())
    }

    async fn bind_udp(&self, bind_ip: IpAddr, port: u16, only_v6: Option<bool>) -> IoResult<UdpSocket> {
        let domain = match bind_ip {
            IpAddr::V4(_) => Domain::IPV4,
            IpAddr::V6(_) => Domain::IPV6,
        };

        let socket = Socket::new(domain, Type::DGRAM, Some(Protocol::UDP))?;

        if let Some(only_v6) = only_v6 {
            socket.set_only_v6(only_v6)?;
        }

        socket.set_nonblocking(true)?;

        socket.bind(&SockAddr::from(
            SocketAddr::new(bind_ip, port)
        ))?;

        UdpSocket::from_std(socket.into())
    }

    pub(crate) async fn bind(&self,
                             bind_ip: IpAddr,
                             bind_ports: BTreeSet<u16>,
                             only_v6: Option<bool>,
                             network: &Network) -> Result<u16, Error> {
        let ports = {
            match network {
                Network::Tcp | Network::Udp => {
                    match bind_ports == *EPHEMERAL_PORTS {
                        // workaround for ephemeral ports, make it chosen by OS.
                        true => vec![0],
                        false => bind_ports.into_iter().collect(),
                    }
                }
                Network::Both => bind_ports.into_iter().collect(),
            }
        };

        // NOTICE: const_btree_len is still unstable, so we have to use this workaround.
        if ports.is_empty() {
            return Err(Error::PortDenied);
        }

        for port in ports {
            let tcp_listener = async {
                if network.tcp() {
                    return Some(self.bind_tcp(bind_ip, port, only_v6).await);
                }

                None
            }.await;

            let udp_socket = async {
                if network.udp() {
                    return Some(self.bind_udp(bind_ip, port, only_v6).await);
                }

                None
            }.await;

            match (tcp_listener, udp_socket) {
                (Some(Ok(tcp_listener)), Some(Ok(udp_socket))) => {
                    self.clone().handle_tcp_listener(tcp_listener).await;
                    let mut udp_sessions = self.udp_sessions.lock().await;
                    *udp_sessions = Some(UdpSessions::new(
                        self.clone(),
                        udp_socket,
                        self.max_packet_size,
                    ));

                    return Ok(port);
                }

                (Some(Ok(tcp_listener)), None) => {
                    let port = tcp_listener.local_addr().unwrap().port(); // Get actual port when port is 0
                    self.clone().handle_tcp_listener(tcp_listener).await;
                    return Ok(port);
                }

                (None, Some(Ok(udp_socket))) => {
                    let port = udp_socket.local_addr().unwrap().port(); // Get actual port when port is 0
                    let mut udp_sessions = self.udp_sessions.lock().await;
                    *udp_sessions = Some(UdpSessions::new(
                        self.clone(),
                        udp_socket,
                        self.max_packet_size,
                    ));

                    return Ok(port);
                }

                _ => {}
            }
        }


        Err(Error::BindFailed)
    }
}