use core::{
    fmt,
    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
    task::{Context, Poll},
    time::Duration,
};

use alloc::{boxed::Box, sync::Arc, vec::Vec};
use async_trait::async_trait;
use libk::{
    block,
    error::Error,
    task::runtime::maybe_timeout,
    vfs::{FileReadiness, Socket},
};
use libk_util::{
    queue::BoundedMpmcQueue,
    sync::spin_rwlock::{IrqSafeRwLock, IrqSafeRwLockReadGuard},
};
use yggdrasil_abi::{
    net::{
        options::{self, UdpSocketOptionVariant},
        MessageHeader, MessageHeaderMut,
    },
    option::OptionValue,
};

use crate::{l3::ip, l4};

use super::SocketTable;

pub struct UdpSocket {
    local: IrqSafeRwLock<Option<SocketAddr>>,
    remote: IrqSafeRwLock<Option<SocketAddr>>,

    options: IrqSafeRwLock<Options>,

    // TODO just place packets here for one less copy?
    receive_queue: BoundedMpmcQueue<(SocketAddr, Vec<u8>)>,
}

struct Options {
    broadcast: bool,
    ttl: u8,
}

static UDP_SOCKETS: IrqSafeRwLock<SocketTable<UdpSocket>> = IrqSafeRwLock::new(SocketTable::new());

impl Default for Options {
    fn default() -> Self {
        Self {
            broadcast: false,
            ttl: 64,
        }
    }
}

impl UdpSocket {
    pub fn new() -> Arc<Self> {
        Arc::new(Self {
            local: IrqSafeRwLock::new(None),
            remote: IrqSafeRwLock::new(None),

            options: IrqSafeRwLock::new(Options::default()),

            receive_queue: BoundedMpmcQueue::new(128),
        })
    }

    pub fn get(local: &SocketAddr) -> Option<Arc<UdpSocket>> {
        UDP_SOCKETS.read().get(local)
    }

    pub fn packet_received(&self, source: SocketAddr, data: &[u8]) -> Result<(), Error> {
        self.receive_queue
            .try_push_back((source, Vec::from(data)))
            .map_err(|_| Error::QueueFull)
    }

    // If address is bound, keep it, if not, bind an ephemeral port
    pub fn ensure_address(self: &Arc<Self>, v6: bool) -> Result<u16, Error> {
        let local = self.local.read();
        if let Some(address) = *local {
            Ok(address.port())
        } else {
            let mut local = IrqSafeRwLockReadGuard::upgrade(local);
            let ip = match v6 {
                true => IpAddr::V4(Ipv4Addr::UNSPECIFIED),
                false => IpAddr::V6(Ipv6Addr::UNSPECIFIED),
            };
            let port = UDP_SOCKETS
                .write()
                .bind_to_ephemeral_port(ip, self.clone())?;
            *local = Some(SocketAddr::new(ip, port));
            Ok(port)
        }
    }

    async fn send(
        &self,
        port: u16,
        payload: &[u8],
        destination: SocketAddr,
    ) -> Result<usize, Error> {
        log::info!("udp: send {} bytes to {destination}", payload.len());
        let options = self.options.read();

        // TODO check that destnation family matches self family
        match (options.broadcast, destination.ip()) {
            // TODO broadcast
            (true, _) => return Err(Error::NotImplemented),
            (false, _) => {
                l4::udp::send(
                    port,
                    destination.ip(),
                    destination.port(),
                    options.ttl,
                    payload,
                )
                .await?;
            }
        }

        Ok(payload.len())
    }

    fn store_packet(
        message: &mut MessageHeaderMut,
        packet: &[u8],
        source: &SocketAddr,
    ) -> Result<usize, Error> {
        log::info!("udp: receive {} bytes from {source}", packet.len());
        if packet.len() > message.payload.len() {
            // TODO check how other OSs handle this
            return Err(Error::BufferTooSmall);
        }
        message.payload[..packet.len()].copy_from_slice(packet);
        ip::store_address(&mut message.source, source)?;
        Ok(packet.len())
    }
}

impl FileReadiness for UdpSocket {
    fn poll_read(&self, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
        self.receive_queue.poll_not_empty(cx).map(Ok)
    }
}

#[async_trait]
impl Socket for UdpSocket {
    fn bind(self: Arc<Self>, local: &[u8]) -> Result<(), Error> {
        let local = ip::load_address(local)?;
        log::info!("udp: bind {local}");
        let mut bound = self.local.write();
        if bound.is_some() {
            return Err(Error::InvalidOperation);
        }
        if local.port() == 0 {
            let port = UDP_SOCKETS
                .write()
                .bind_to_ephemeral_port(local.ip(), self.clone())?;
            *bound = Some(SocketAddr::new(local.ip(), port));
        } else {
            UDP_SOCKETS.write().try_insert(local, self.clone())?;
            *bound = Some(local);
        }
        Ok(())
    }

    async fn connect(
        self: Arc<Self>,
        remote: &[u8],
        _timeout: Option<Duration>,
    ) -> Result<(), Error> {
        let remote = ip::load_address(remote)?;
        let mut connected = self.remote.write();
        if connected.is_some() {
            return Err(Error::InvalidOperation);
        }
        *connected = Some(remote);
        Ok(())
    }

    fn close(self: Arc<Self>) -> Result<(), Error> {
        if let Some(local) = self.local.write().take() {
            log::debug!("UDP socket closed: {}", local);
            UDP_SOCKETS.write().remove(local)?;
        }
        Ok(())
    }

    fn set_option(&self, option: u32, buffer: &[u8]) -> Result<(), Error> {
        let option = UdpSocketOptionVariant::try_from(option)?;
        let mut options = self.options.write();
        match option {
            UdpSocketOptionVariant::Ttl => {
                let ttl = options::Ttl::load(buffer)?;
                if !(1..256).contains(&ttl) {
                    return Err(Error::InvalidArgument);
                }
                options.ttl = ttl as u8;
                Ok(())
            }
            UdpSocketOptionVariant::Broadcast => {
                options.broadcast = options::Broadcast::load(buffer)?;
                Ok(())
            }
            UdpSocketOptionVariant::MulticastTtlV4
            | UdpSocketOptionVariant::MulticastLoopV4
            | UdpSocketOptionVariant::MulticastLoopV6 => {
                log::warn!("TODO: udp multicast not yet implemented");
                Err(Error::InvalidArgument)
            }
        }
    }

    fn get_option(&self, option: u32, buffer: &mut [u8]) -> Result<usize, Error> {
        let option = UdpSocketOptionVariant::try_from(option)?;
        let options = self.options.read();
        match option {
            UdpSocketOptionVariant::Ttl => {
                let ttl = options.ttl as u32;
                options::Ttl::store(&ttl, buffer)
            }
            UdpSocketOptionVariant::Broadcast => {
                options::Broadcast::store(&options.broadcast, buffer)
            }
            UdpSocketOptionVariant::MulticastLoopV6
            | UdpSocketOptionVariant::MulticastLoopV4
            | UdpSocketOptionVariant::MulticastTtlV4 => {
                log::warn!("TODO: udp multicast not yet implemented");
                Err(Error::InvalidArgument)
            }
        }
    }

    fn send_nonblocking(self: Arc<Self>, message: &MessageHeader) -> Result<usize, Error> {
        log::warn!("TODO: UDP::send_nonblocking()");
        block!(self.send_message(message, None).await)?
    }
    async fn send_message(
        self: Arc<Self>,
        message: &MessageHeader,
        timeout: Option<Duration>,
    ) -> Result<usize, Error> {
        let _ = timeout;
        let destination = message
            .destination
            .map(ip::load_address)
            .transpose()?
            .or_else(|| *self.remote.read())
            .ok_or(Error::NotConnected)?;

        // If socket wasn't bound yet, bind it to an ephemeral port
        let port = self.ensure_address(destination.ip().is_ipv6())?;

        self.send(port, message.payload, destination).await
    }

    fn receive_nonblocking(
        self: Arc<Self>,
        message: &mut MessageHeaderMut,
    ) -> Result<usize, Error> {
        let (source, packet) = self
            .receive_queue
            .try_pop_front()
            .ok_or(Error::WouldBlock)?;
        Self::store_packet(message, &packet, &source)
    }
    async fn receive_message(
        self: Arc<Self>,
        message: &mut MessageHeaderMut,
        timeout: Option<Duration>,
    ) -> Result<usize, Error> {
        let future = self.receive_queue.pop_front();
        let (source, packet) = maybe_timeout(future, timeout).await?;
        Self::store_packet(message, &packet, &source)
    }
}

impl fmt::Debug for UdpSocket {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let local = *self.local.read();
        let remote = *self.remote.read();

        f.debug_struct("UdpSocket")
            .field("local", &local)
            .field("remote", &remote)
            .finish_non_exhaustive()
    }
}
