use core::{
    fmt,
    future::poll_fn,
    net::SocketAddr,
    task::{Context, Poll},
    time::Duration,
};

use alloc::sync::Arc;
use libk::{
    error::Error,
    task::runtime::{maybe_timeout, with_timeout},
    time::monotonic_time,
};
use libk_util::sync::spin_rwlock::{IrqSafeRwLock, IrqSafeRwLockWriteGuard};

use crate::{
    interface::NetworkInterface,
    l3::Route,
    l4::tcp::{TcpConnection, TcpConnectionState},
    socket::TwoWaySocketTable,
};

use super::TcpListener;

pub struct TcpStream {
    pub(super) local: SocketAddr,
    pub(super) remote: SocketAddr,

    // Listener which accepted the socket
    listener: Option<Arc<TcpListener>>,
    pub(crate) connection: IrqSafeRwLock<TcpConnection>,
}

static TCP_STREAMS: IrqSafeRwLock<TwoWaySocketTable<TcpStream>> =
    IrqSafeRwLock::new(TwoWaySocketTable::new());

impl TcpStream {
    pub async fn connect(
        remote: SocketAddr,
        timeout: Option<Duration>,
    ) -> Result<Arc<Self>, Error> {
        log::info!("tcp: try connect to {remote} (timeout={timeout:?})");
        // Lookup route to remote
        let (interface_id, _, remote_ip) =
            Route::lookup(remote.ip()).ok_or(Error::HostUnreachable)?;
        let remote = SocketAddr::new(remote_ip, remote.port());
        let interface = NetworkInterface::get(interface_id)?;
        let local_ip = interface.address.read().ok_or(Error::NetworkUnreachable)?;

        // Create a new TcpStream with an ephemeral port
        let stream = {
            let mut streams = TCP_STREAMS.write();
            streams.try_insert_with_ephemeral_port(local_ip, remote, |port| {
                let t = monotonic_time();
                let tx_seq = t.as_millis() as u32;
                let local = SocketAddr::new(local_ip, port);
                let connection =
                    TcpConnection::new(local, remote, 16384, tx_seq, 0, TcpConnectionState::Closed);

                Ok(Arc::new(Self {
                    local,
                    remote,
                    listener: None,
                    connection: IrqSafeRwLock::new(connection),
                }))
            })?
        };

        let deadline = timeout.map(|t| monotonic_time() + t);
        // TODO limit interval between attempts to timeout value if one is provided
        let mut t = 200;
        for i in 0..5 {
            let timeout = Duration::from_millis(t);
            log::debug!("Try SYN with timeout={:?}", timeout);
            match stream.try_connect(timeout).await {
                Ok(()) => return Ok(stream),
                Err(Error::TimedOut) => {
                    log::warn!("tcp: connect attempt [{}/5] to {remote} timed out", i + 1);

                    if let Some(deadline) = deadline
                        && monotonic_time() > deadline
                    {
                        break;
                    }
                }
                Err(error) => return Err(error),
            }
            t *= 2;
        }

        // Couldn't establish
        stream.close(false).await.ok();

        log::warn!("tcp: connection to {remote} timed out");

        Err(Error::TimedOut)
    }

    async fn try_connect(&self, timeout: Duration) -> Result<(), Error> {
        {
            let mut connection = self.connection.write();
            connection.send_syn().await?;
        }

        let fut = poll_fn(|cx| {
            let connection = self.connection.read();
            connection.poll_established(cx)
        });

        with_timeout(fut, timeout).await?
    }

    pub fn accept(self: &Arc<Self>) {
        if let Some(listener) = self.listener.as_ref() {
            listener.accept_stream(self.clone());
        }
    }

    pub fn accept_remote(
        listener: Arc<TcpListener>,
        local: SocketAddr,
        remote: SocketAddr,
        remote_window_size: usize,
        tx_seq: u32,
        rx_seq: u32,
    ) -> Result<Arc<TcpStream>, Error> {
        let mut streams = TCP_STREAMS.write();

        streams.try_insert_with(local, remote, move || {
            let connection = TcpConnection::new(
                local,
                remote,
                remote_window_size,
                tx_seq,
                rx_seq,
                TcpConnectionState::SynReceived,
            );

            log::debug!("Accepted TCP socket {} -> {}", local, remote);

            let socket = Self {
                local,
                remote,
                listener: Some(listener),
                connection: IrqSafeRwLock::new(connection),
            };

            Ok(Arc::new(socket))
        })
    }

    pub async fn close(&self, remove_from_listener: bool) -> Result<(), Error> {
        // TODO timeout here
        // Already closing
        if self.connection.read().is_closing() {
            return Ok(());
        }

        // Wait for all sent data to be acknowledged
        {
            let mut connection = poll_fn(|cx| {
                let connection = self.connection.write();
                match connection.poll_send(cx) {
                    Poll::Ready(Ok(())) => Poll::Ready(Ok(connection)),
                    Poll::Ready(Err(error)) => Poll::Ready(Err(error)),
                    Poll::Pending => Poll::Pending,
                }
            })
            .await?;

            connection.finish().await?;
        }

        log::debug!(
            "TCP socket closed (FinWait2/Closed): {} <-> {}",
            self.local,
            self.remote
        );

        // Wait for connection to get closed
        poll_fn(|cx| {
            let connection = self.connection.read();
            connection.poll_finish(cx)
        })
        .await;

        if remove_from_listener {
            if let Some(listener) = self.listener.as_ref() {
                listener.remove_stream(self.remote);
            };
        }

        Ok(())
    }

    pub(crate) fn remove_stream(&self) -> Result<(), Error> {
        log::debug!(
            "TCP socket closed and removed: {} <-> {}",
            self.local,
            self.remote
        );
        let connection = self.connection.read();
        debug_assert!(connection.is_closed());
        TCP_STREAMS.write().remove(self.local, self.remote)?;
        connection.notify_all();
        Ok(())
    }

    pub fn get(local: SocketAddr, remote: SocketAddr) -> Option<Arc<TcpStream>> {
        TCP_STREAMS.read().get(local, remote)
    }

    pub(super) fn poll_receive(
        &self,
        cx: &mut Context<'_>,
    ) -> Poll<Result<IrqSafeRwLockWriteGuard<'_, TcpConnection>, Error>> {
        let lock = self.connection.write();
        match lock.poll_receive(cx) {
            Poll::Ready(Ok(())) => Poll::Ready(Ok(lock)),
            Poll::Ready(Err(error)) => Poll::Ready(Err(error)),
            Poll::Pending => Poll::Pending,
        }
    }

    pub async fn receive(
        &self,
        buffer: &mut [u8],
        timeout: Option<Duration>,
    ) -> Result<usize, Error> {
        let future = poll_fn(|cx| match self.poll_receive(cx) {
            Poll::Ready(Ok(mut lock)) => Poll::Ready(lock.read_nonblocking(buffer)),
            Poll::Ready(Err(error)) => Poll::Ready(Err(error)),
            Poll::Pending => Poll::Pending,
        });

        maybe_timeout(future, timeout).await.flatten()
    }

    pub async fn send(&self, data: &[u8]) -> Result<usize, Error> {
        let mut pos = 0;
        let mut rem = data.len();
        while rem != 0 {
            // TODO check MTU
            let amount = rem.min(512);
            self.send_segment(&data[pos..pos + amount]).await?;
            pos += amount;
            rem -= amount;
        }
        Ok(pos)
    }

    pub fn receive_nonblocking(&self, buffer: &mut [u8]) -> Result<usize, Error> {
        if buffer.is_empty() {
            return Ok(0);
        }
        let mut lock = self.connection.write();
        match lock.read_nonblocking(buffer) {
            Ok(0) => Err(Error::WouldBlock),
            res => res,
        }
    }

    pub fn send_nonblocking(&self, data: &[u8]) -> Result<usize, Error> {
        if data.is_empty() {
            return Ok(0);
        }
        let mut pos = 0;
        let mut rem = data.len();
        let mut sent = false;
        while rem != 0 {
            // TODO check MTU
            let amount = rem.min(512);
            match self.send_segment_nonblocking(&data[pos..pos + amount]) {
                Ok(()) => sent = true,
                Err(Error::WouldBlock) => break,
                Err(error) => return Err(error),
            }
            pos += amount;
            rem -= amount;
        }
        // No data sent and WouldBlock returned
        if !sent {
            return Err(Error::WouldBlock);
        }
        Ok(pos)
    }

    async fn send_segment(&self, data: &[u8]) -> Result<(), Error> {
        {
            let mut connection = poll_fn(|cx| {
                let connection = self.connection.write();
                match connection.poll_send(cx) {
                    Poll::Ready(Ok(())) => Poll::Ready(Ok(connection)),
                    Poll::Ready(Err(error)) => Poll::Ready(Err(error)),
                    Poll::Pending => Poll::Pending,
                }
            })
            .await?;

            connection.transmit(data).await?;
        }

        poll_fn(|cx| {
            let connection = self.connection.read();
            connection.poll_acknowledge(cx)
        })
        .await;

        Ok(())
    }

    fn send_segment_nonblocking(&self, _data: &[u8]) -> Result<(), Error> {
        todo!()
    }
}

impl fmt::Debug for TcpStream {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("TcpStream")
            .field("local", &self.local)
            .field("remote", &self.remote)
            .finish_non_exhaustive()
    }
}
