use core::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6};

use alloc::{
    collections::{btree_map::Entry, BTreeMap},
    sync::Arc,
};
use yggdrasil_abi::error::Error;

pub mod udp;
pub use udp::UdpSocket;
pub mod tcp;
pub use tcp::TcpSocket;
pub mod raw;
pub use raw::RawSocket;
pub mod config;
pub use config::NetConfigSocket;
pub mod local;
pub use local::LocalPacketSocket;

pub struct SocketTable<T> {
    inner: BTreeMap<SocketAddr, Arc<T>>,
}

pub struct TwoWaySocketTable<T> {
    inner: BTreeMap<(SocketAddr, SocketAddr), Arc<T>>,
    last_port: u16,
}

impl<T> TwoWaySocketTable<T> {
    pub const fn new() -> Self {
        Self {
            inner: BTreeMap::new(),
            last_port: 32768,
        }
    }

    pub fn try_insert_with<F: FnOnce() -> Result<Arc<T>, Error>>(
        &mut self,
        local: SocketAddr,
        remote: SocketAddr,
        with: F,
    ) -> Result<Arc<T>, Error> {
        if self.inner.contains_key(&(local, remote)) {
            return Err(Error::AddrInUse);
        }
        let socket = with()?;
        self.inner.insert((local, remote), socket.clone());
        Ok(socket)
    }

    const fn next_port_wrapping(port: u16) -> u16 {
        let port = port + 1;
        if port >= u16::MAX - 1 {
            32768
        } else {
            port
        }
    }

    fn try_insert_into_range<F: FnMut(u16) -> Result<Arc<T>, Error>, R: Iterator<Item = u16>>(
        &mut self,
        local: IpAddr,
        remote: SocketAddr,
        range: R,
        with: &mut F,
    ) -> Result<(Arc<T>, u16), Error> {
        for port in range {
            let local = SocketAddr::new(local, port);

            match self.try_insert_with(local, remote, || with(port)) {
                Ok(socket) => return Ok((socket, port)),
                Err(Error::AddrInUse) => continue,
                Err(error) => return Err(error),
            }
        }

        Err(Error::AddrInUse)
    }

    pub fn try_insert_with_ephemeral_port<F: FnMut(u16) -> Result<Arc<T>, Error>>(
        &mut self,
        local: IpAddr,
        remote: SocketAddr,
        mut with: F,
    ) -> Result<Arc<T>, Error> {
        if let Ok((socket, port)) =
            self.try_insert_into_range(local, remote, self.last_port..u16::MAX - 1, &mut with)
        {
            self.last_port = Self::next_port_wrapping(port);
            return Ok(socket);
        }
        if let Ok((socket, port)) =
            self.try_insert_into_range(local, remote, 32768..self.last_port, &mut with)
        {
            self.last_port = Self::next_port_wrapping(port);
            return Ok(socket);
        }

        Err(Error::AddrInUse)
    }

    pub fn remove(&mut self, local: SocketAddr, remote: SocketAddr) -> Result<(), Error> {
        match self.inner.remove(&(local, remote)) {
            Some(_) => Ok(()),
            None => Err(Error::DoesNotExist),
        }
    }

    pub fn get(&self, local: SocketAddr, remote: SocketAddr) -> Option<Arc<T>> {
        self.inner.get(&(local, remote)).cloned()
    }
}

impl<T> SocketTable<T> {
    pub const fn new() -> Self {
        Self {
            inner: BTreeMap::new(),
        }
    }

    pub fn try_insert_with_ephemeral_port<F: FnMut(u16) -> Result<Arc<T>, Error>>(
        &mut self,
        local: IpAddr,
        mut with: F,
    ) -> Result<Arc<T>, Error> {
        for port in 32768..u16::MAX - 1 {
            let local = SocketAddr::new(local, port);

            match self.try_insert_with(local, || with(port)) {
                Ok(socket) => return Ok(socket),
                Err(Error::AddrInUse) => continue,
                Err(error) => return Err(error),
            }
        }
        Err(Error::AddrInUse)
    }

    pub fn bind_to_ephemeral_port(&mut self, ip: IpAddr, socket: Arc<T>) -> Result<u16, Error> {
        for port in 32768..u16::MAX - 1 {
            let local = SocketAddr::new(ip, port);

            match self.try_insert(local, socket.clone()) {
                Ok(()) => return Ok(port),
                Err(Error::AddrInUse) => continue,
                Err(error) => return Err(error),
            }
        }
        Err(Error::AddrInUse)
    }

    pub fn try_insert(&mut self, address: SocketAddr, socket: Arc<T>) -> Result<(), Error> {
        match self.inner.entry(address) {
            Entry::Vacant(entry) => {
                entry.insert(socket);
                Ok(())
            }
            Entry::Occupied(_) => Err(Error::AddrInUse),
        }
    }

    pub fn try_insert_with<F: FnOnce() -> Result<Arc<T>, Error>>(
        &mut self,
        address: SocketAddr,
        with: F,
    ) -> Result<Arc<T>, Error> {
        if self.inner.contains_key(&address) {
            return Err(Error::AddrInUse);
        }
        let socket = with()?;
        self.inner.insert(address, socket.clone());
        Ok(socket)
    }

    pub fn remove(&mut self, local: SocketAddr) -> Result<(), Error> {
        match self.inner.remove(&local) {
            Some(_) => Ok(()),
            None => Err(Error::DoesNotExist),
        }
    }

    pub fn get_exact(&self, local: &SocketAddr) -> Option<Arc<T>> {
        self.inner.get(local).cloned()
    }

    pub fn get(&self, local: &SocketAddr) -> Option<Arc<T>> {
        if let Some(socket) = self.inner.get(local) {
            return Some(socket.clone());
        }

        let unspec = match local {
            SocketAddr::V4(_v4) => SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, local.port()).into(),
            SocketAddr::V6(_v6) => {
                SocketAddrV6::new(Ipv6Addr::UNSPECIFIED, local.port(), 0, 0).into()
            }
        };

        self.inner.get(&unspec).cloned()
    }
}
