use core::{
    fmt, mem,
    sync::atomic::Ordering,
    task::{Context, Poll},
    time::Duration,
};

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

mod listener;
mod stream;

pub use listener::TcpListener;
pub use stream::TcpStream;

use crate::l3::ip;

pub struct TcpSocket {
    options: IrqSafeRwLock<TcpSocketOptions>,
    inner: IrqSafeRwLock<TcpSocketInner>,
}

pub enum TcpSocketInner {
    Empty,
    Listener(Arc<TcpListener>),
    Stream(Arc<TcpStream>),
}

struct TcpSocketOptions {
    nodelay: bool,
    v6_only: bool,
}

#[allow(clippy::derivable_impls)]
impl Default for TcpSocketOptions {
    fn default() -> Self {
        Self {
            nodelay: false,
            v6_only: false,
        }
    }
}

impl TcpSocket {
    pub fn new() -> Arc<Self> {
        Arc::new(Self {
            options: IrqSafeRwLock::new(TcpSocketOptions::default()),
            inner: IrqSafeRwLock::new(TcpSocketInner::Empty),
        })
    }

    fn as_stream(&self) -> Option<Arc<TcpStream>> {
        if let TcpSocketInner::Stream(stream) = &*self.inner.read() {
            Some(stream.clone())
        } else {
            None
        }
    }

    fn as_listener(&self) -> Option<Arc<TcpListener>> {
        if let TcpSocketInner::Listener(listener) = &*self.inner.read() {
            Some(listener.clone())
        } else {
            None
        }
    }
}

#[async_trait]
impl Socket for TcpSocket {
    fn bind(self: Arc<Self>, local: &[u8]) -> Result<(), Error> {
        let local = ip::load_address(local)?;
        let mut inner = self.inner.write();

        // Already connected or bound
        if !matches!(&*inner, TcpSocketInner::Empty) {
            return Err(Error::InvalidOperation);
        }

        // Bind a listener socket
        let listener = TcpListener::bind(local)?;
        *inner = TcpSocketInner::Listener(listener);

        Ok(())
    }

    fn listen(self: Arc<Self>) -> Result<(), Error> {
        let listener = self.as_listener().ok_or(Error::InvalidOperation)?;
        match listener
            .listening
            .compare_exchange(false, true, Ordering::Release, Ordering::Relaxed)
        {
            Ok(_) => Ok(()),
            Err(_) => Err(Error::InvalidOperation),
        }
    }

    async fn accept(
        self: Arc<Self>,
        mut address: Option<&mut [u8]>,
        timeout: Option<Duration>,
    ) -> Result<Arc<dyn Socket>, Error> {
        let _ = timeout;
        let listener = self.as_listener().ok_or(Error::InvalidOperation)?;
        let stream = listener.accept().await?;
        let remote = stream.remote;
        let socket = Arc::new(TcpSocket {
            options: IrqSafeRwLock::new(TcpSocketOptions::default()),
            inner: IrqSafeRwLock::new(TcpSocketInner::Stream(stream)),
        });
        // TODO if this store fails, socket leaks
        ip::store_address(&mut address, &remote)?;
        Ok(socket)
    }
    fn accept_nonblocking(
        self: Arc<Self>,
        mut address: Option<&mut [u8]>,
    ) -> Result<Arc<dyn Socket>, Error> {
        let listener = self.as_listener().ok_or(Error::InvalidOperation)?;
        let stream = listener.accept_nonblocking()?;
        let remote = stream.remote;
        let socket = Arc::new(TcpSocket {
            options: IrqSafeRwLock::new(TcpSocketOptions::default()),
            inner: IrqSafeRwLock::new(TcpSocketInner::Stream(stream)),
        });
        // TODO if this store fails, socket leaks
        ip::store_address(&mut address, &remote)?;
        Ok(socket)
    }

    async fn connect(
        self: Arc<Self>,
        remote: &[u8],
        timeout: Option<Duration>,
    ) -> Result<(), Error> {
        let remote = ip::load_address(remote)?;
        let mut inner = self.inner.write();

        // Already connected or bound
        if !matches!(&*inner, TcpSocketInner::Empty) {
            return Err(Error::InvalidOperation);
        }

        let stream = TcpStream::connect(remote, timeout).await?;
        *inner = TcpSocketInner::Stream(stream);

        Ok(())
    }

    fn close(self: Arc<Self>) -> Result<(), Error> {
        let inner = mem::replace(&mut *self.inner.write(), TcpSocketInner::Empty);

        match inner {
            TcpSocketInner::Empty => Ok(()),
            TcpSocketInner::Stream(socket) => block!(socket.close(true).await)?,
            TcpSocketInner::Listener(socket) => socket.close(),
        }
    }

    async fn shutdown(self: Arc<Self>, read: bool, write: bool) -> Result<(), Error> {
        log::warn!("TODO: shutdown(read = {read}, write = {write})");
        Ok(())
    }

    // fn local_address(&self) -> Option<SocketAddr> {
    //     match &*self.inner.read() {
    //         TcpSocketInner::Empty => None,
    //         TcpSocketInner::Stream(socket) => Some(socket.local),
    //         TcpSocketInner::Listener(socket) => Some(socket.local),
    //     }
    // }

    // fn remote_address(&self) -> Option<SocketAddr> {
    //     match &*self.inner.read() {
    //         TcpSocketInner::Empty | TcpSocketInner::Listener(_) => None,
    //         TcpSocketInner::Stream(socket) => Some(socket.remote),
    //     }
    // }

    fn set_option(&self, option: u32, buffer: &[u8]) -> Result<(), Error> {
        let option = TcpSocketOptionVariant::try_from(option)?;
        let mut options = self.options.write();
        match option {
            TcpSocketOptionVariant::NoDelay => {
                options.nodelay = options::NoDelay::load(buffer)?;
                Ok(())
            }
            TcpSocketOptionVariant::Ipv6Only => {
                options.v6_only = options::Ipv6Only::load(buffer)?;
                Ok(())
            }
        }
    }

    fn get_option(&self, option: u32, buffer: &mut [u8]) -> Result<usize, Error> {
        let option = TcpSocketOptionVariant::try_from(option)?;
        let options = self.options.read();
        match option {
            TcpSocketOptionVariant::NoDelay => options::NoDelay::store(&options.nodelay, buffer),
            TcpSocketOptionVariant::Ipv6Only => options::Ipv6Only::store(&options.v6_only, buffer),
        }
    }

    fn send_nonblocking(self: Arc<Self>, message: &MessageHeader) -> Result<usize, Error> {
        let stream = self.as_stream().ok_or(Error::NotConnected)?;
        let len = stream.send_nonblocking(message.payload)?;
        Ok(len)
    }
    async fn send_message(
        self: Arc<Self>,
        message: &MessageHeader,
        timeout: Option<Duration>,
    ) -> Result<usize, Error> {
        let stream = self.as_stream().ok_or(Error::NotConnected)?;
        let future = stream.send(message.payload);
        maybe_timeout(future, timeout).await?
    }

    fn receive_nonblocking(
        self: Arc<Self>,
        message: &mut MessageHeaderMut,
    ) -> Result<usize, Error> {
        let stream = self.as_stream().ok_or(Error::NotConnected)?;
        let len = stream.receive_nonblocking(message.payload)?;
        ip::store_address(&mut message.source, &stream.remote)?;
        Ok(len)
    }
    async fn receive_message(
        self: Arc<Self>,
        message: &mut MessageHeaderMut,
        timeout: Option<Duration>,
    ) -> Result<usize, Error> {
        let stream = self.as_stream().ok_or(Error::NotConnected)?;
        let len = stream.receive(message.payload, timeout).await?;
        ip::store_address(&mut message.source, &stream.remote)?;
        Ok(len)
    }
}

impl FileReadiness for TcpSocket {
    fn poll_read(&self, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
        match &*self.inner.read() {
            TcpSocketInner::Empty => Poll::Ready(Ok(())),
            TcpSocketInner::Stream(socket) => socket.poll_receive(cx).map_ok(|_| ()),
            TcpSocketInner::Listener(socket) => socket.poll_accept(cx).map(|_| Ok(())),
        }
    }
}

impl fmt::Debug for TcpSocket {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match &*self.inner.read() {
            TcpSocketInner::Empty => f.debug_struct("TcpSocket").finish_non_exhaustive(),
            TcpSocketInner::Stream(socket) => f
                .debug_struct("TcpSocket")
                .field("stream", socket)
                .finish_non_exhaustive(),
            TcpSocketInner::Listener(socket) => f
                .debug_struct("TcpSocket")
                .field("listener", socket)
                .finish_non_exhaustive(),
        }
    }
}
