use core::{
    fmt,
    task::{Context, Poll},
    time::Duration,
};

use alloc::{boxed::Box, sync::Arc};
use async_trait::async_trait;
use libk_util::sync::spin_rwlock::IrqSafeRwLock;
use yggdrasil_abi::{
    error::Error,
    net::{
        options::{self, SocketOptionVariant},
        MessageHeader, MessageHeaderMut, SocketShutdown,
    },
    option::OptionValue,
};

use crate::vfs::FileReadiness;

struct InnerOptions {
    connect_timeout: Option<Duration>,
    recv_timeout: Option<Duration>,
    send_timeout: Option<Duration>,
}

pub struct SocketWrapper {
    inner: Arc<dyn Socket>,
    options: IrqSafeRwLock<InnerOptions>,
}

/// Interface for interacting with network sockets
#[async_trait]
pub trait Socket: FileReadiness + fmt::Debug + Send {
    fn bind(self: Arc<Self>, local: &[u8]) -> Result<(), Error> {
        let _ = local;
        Err(Error::InvalidOperation)
    }
    async fn connect(
        self: Arc<Self>,
        remote: &[u8],
        timeout: Option<Duration>,
    ) -> Result<(), Error>;

    fn listen(self: Arc<Self>) -> Result<(), Error> {
        Err(Error::InvalidOperation)
    }
    async fn accept(
        self: Arc<Self>,
        address: Option<&mut [u8]>,
        timeout: Option<Duration>,
    ) -> Result<Arc<dyn Socket>, Error> {
        let _ = address;
        let _ = timeout;
        Err(Error::InvalidOperation)
    }
    fn accept_nonblocking(
        self: Arc<Self>,
        address: Option<&mut [u8]>,
    ) -> Result<Arc<dyn Socket>, Error> {
        let _ = address;
        Err(Error::InvalidOperation)
    }

    async fn shutdown(self: Arc<Self>, read: bool, write: bool) -> Result<(), Error> {
        let _ = read;
        let _ = write;
        Err(Error::InvalidOperation)
    }

    /// Closes a socket
    fn close(self: Arc<Self>) -> Result<(), Error> {
        Ok(())
    }

    /// Updates a socket option
    fn set_option(&self, option: u32, buffer: &[u8]) -> Result<(), Error> {
        let _ = option;
        let _ = buffer;
        Err(Error::InvalidOperation)
    }

    /// Gets a socket option
    fn get_option(&self, option: u32, buffer: &mut [u8]) -> Result<usize, Error> {
        let _ = option;
        let _ = buffer;
        Err(Error::InvalidOperation)
    }

    fn send_nonblocking(self: Arc<Self>, message: &MessageHeader) -> Result<usize, Error>;
    async fn send_message(
        self: Arc<Self>,
        message: &MessageHeader,
        timeout: Option<Duration>,
    ) -> Result<usize, Error>;

    fn receive_nonblocking(self: Arc<Self>, message: &mut MessageHeaderMut)
        -> Result<usize, Error>;
    async fn receive_message(
        self: Arc<Self>,
        message: &mut MessageHeaderMut,
        timeout: Option<Duration>,
    ) -> Result<usize, Error>;
}

impl SocketWrapper {
    pub fn from_inner(inner: Arc<dyn Socket>) -> Self {
        Self::from(inner)
    }

    pub fn connect(&self, remote: &[u8]) -> Result<(), Error> {
        let timeout = self.options.read().connect_timeout;
        block!(self.inner.clone().connect(remote, timeout).await)?
    }

    pub fn accept(&self, remote: Option<&mut [u8]>, non_blocking: bool) -> Result<Self, Error> {
        let timeout = self.options.read().recv_timeout;
        let socket = if non_blocking {
            self.inner.clone().accept_nonblocking(remote)?
        } else {
            block!(self.inner.clone().accept(remote, timeout).await)??
        };
        Ok(Self::from_inner(socket))
    }

    pub fn send_message(
        &self,
        message: &MessageHeader,
        non_blocking: bool,
    ) -> Result<usize, Error> {
        if non_blocking {
            self.inner.clone().send_nonblocking(message)
        } else {
            let timeout = self.options.read().send_timeout;
            block!(self.inner.clone().send_message(message, timeout).await)?
        }
    }

    pub fn receive_message(
        &self,
        message: &mut MessageHeaderMut,
        non_blocking: bool,
    ) -> Result<usize, Error> {
        if non_blocking {
            self.inner.clone().receive_nonblocking(message)
        } else {
            let timeout = self.options.read().recv_timeout;
            block!(self.inner.clone().receive_message(message, timeout).await)?
        }
    }

    pub fn bind(&self, local: &[u8]) -> Result<(), Error> {
        self.inner.clone().bind(local)
    }

    pub fn listen(&self) -> Result<(), Error> {
        self.inner.clone().listen()
    }

    pub fn shutdown(&self, how: SocketShutdown) -> Result<(), Error> {
        block!(
            self.inner
                .clone()
                .shutdown(
                    how.contains(SocketShutdown::READ),
                    how.contains(SocketShutdown::WRITE)
                )
                .await
        )?
    }

    pub fn set_option(&self, option: u32, buffer: &[u8]) -> Result<(), Error> {
        if let Ok(option) = SocketOptionVariant::try_from(option) {
            let mut options = self.options.write();
            return match option {
                SocketOptionVariant::RecvTimeout => {
                    options.recv_timeout = options::RecvTimeout::load(buffer)?;
                    Ok(())
                }
                SocketOptionVariant::SendTimeout => {
                    options.send_timeout = options::SendTimeout::load(buffer)?;
                    Ok(())
                }
                SocketOptionVariant::ConnectTimeout => {
                    options.connect_timeout = options::ConnectTimeout::load(buffer)?;
                    Ok(())
                }
                SocketOptionVariant::LocalAddress | SocketOptionVariant::PeerAddress => {
                    Err(Error::InvalidArgument)
                }
            };
        }

        // Not a socket-level option, pass to the specific protocol level
        self.inner.set_option(option, buffer)
    }

    pub fn get_option(&self, option: u32, buffer: &mut [u8]) -> Result<usize, Error> {
        if let Ok(option) = SocketOptionVariant::try_from(option) {
            let options = self.options.read();
            return match option {
                SocketOptionVariant::LocalAddress => {
                    todo!()
                    // let address = self.local_address().map(Into::into);
                    // options::LocalAddress::store(&address, buffer)
                }
                SocketOptionVariant::PeerAddress => {
                    todo!()
                    // let address = self.local_address().map(Into::into);
                    // options::PeerAddress::store(&address, buffer)
                }
                SocketOptionVariant::RecvTimeout => {
                    options::RecvTimeout::store(&options.recv_timeout, buffer)
                }
                SocketOptionVariant::SendTimeout => {
                    options::SendTimeout::store(&options.send_timeout, buffer)
                }
                SocketOptionVariant::ConnectTimeout => {
                    options::ConnectTimeout::store(&options.connect_timeout, buffer)
                }
            };
        }

        // Not a socket-level option, pass to the specific protocol level
        self.inner.get_option(option, buffer)
    }
}

impl From<Arc<dyn Socket>> for SocketWrapper {
    fn from(value: Arc<dyn Socket>) -> Self {
        Self {
            inner: value,
            options: IrqSafeRwLock::new(InnerOptions::default()),
        }
    }
}

impl FileReadiness for SocketWrapper {
    fn poll_read(&self, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
        self.inner.poll_read(cx)
    }
}

impl fmt::Debug for SocketWrapper {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Debug::fmt(self.inner.as_ref(), f)
    }
}

impl Drop for SocketWrapper {
    fn drop(&mut self) {
        let res = self.inner.clone().close();
        if let Err(error) = res {
            log::warn!("Socket close error: {error:?}");
        }
    }
}

impl Default for InnerOptions {
    fn default() -> Self {
        Self {
            connect_timeout: None,
            recv_timeout: None,
            send_timeout: None,
        }
    }
}
