use core::{
    fmt,
    future::poll_fn,
    net::SocketAddr,
    sync::atomic::{AtomicBool, Ordering},
    task::{Context, Poll},
};

use alloc::{collections::btree_map::BTreeMap, sync::Arc, vec::Vec};
use libk::error::Error;
use libk_util::{
    sync::{spin_rwlock::IrqSafeRwLock, IrqSafeSpinlock, IrqSafeSpinlockGuard},
    waker::QueueWaker,
};

use crate::socket::SocketTable;

use super::TcpStream;

pub struct TcpListener {
    pub(super) local: SocketAddr,
    pub(super) listening: AtomicBool,

    // Currently active sockets
    sockets: IrqSafeRwLock<BTreeMap<SocketAddr, Arc<TcpStream>>>,
    pending_accept: IrqSafeSpinlock<Vec<Arc<TcpStream>>>,
    accept_notify: QueueWaker,
}

static TCP_LISTENERS: IrqSafeRwLock<SocketTable<TcpListener>> =
    IrqSafeRwLock::new(SocketTable::new());

impl TcpListener {
    pub fn bind(local: SocketAddr) -> Result<Arc<Self>, Error> {
        if local.port() == 0 {
            // TODO: ephemeral binding
            todo!();
        }

        let listener = Arc::new(TcpListener {
            local,
            listening: AtomicBool::new(false),

            sockets: IrqSafeRwLock::new(BTreeMap::new()),
            pending_accept: IrqSafeSpinlock::new(Vec::new()),
            accept_notify: QueueWaker::new(),
        });
        TCP_LISTENERS.write().try_insert(local, listener.clone())?;

        Ok(listener)
    }

    pub fn close(&self) -> Result<(), Error> {
        // TODO if clients not closed already, send RST?
        TCP_LISTENERS.write().remove(self.local)
    }

    pub(super) async fn accept(&self) -> Result<Arc<TcpStream>, Error> {
        poll_fn(|cx| match self.poll_accept(cx) {
            Poll::Ready(mut lock) => Poll::Ready(Ok(lock.pop().unwrap())),
            Poll::Pending => Poll::Pending,
        })
        .await
    }

    pub(super) fn accept_nonblocking(&self) -> Result<Arc<TcpStream>, Error> {
        self.pending_accept.lock().pop().ok_or(Error::WouldBlock)
    }

    pub(super) fn poll_accept(
        &self,
        cx: &mut Context<'_>,
    ) -> Poll<IrqSafeSpinlockGuard<'_, Vec<Arc<TcpStream>>>> {
        let lock = self.pending_accept.lock();
        self.accept_notify.register(cx.waker());
        if !lock.is_empty() {
            self.accept_notify.remove(cx.waker());
            Poll::Ready(lock)
        } else {
            Poll::Pending
        }
    }

    pub(super) fn accept_stream(&self, stream: Arc<TcpStream>) {
        log::debug!("{}: accept {}", self.local, stream.remote);
        self.sockets.write().insert(stream.remote, stream.clone());
        self.pending_accept.lock().push(stream);
        self.accept_notify.wake_one();
    }

    pub(super) fn remove_stream(&self, remote: SocketAddr) {
        log::debug!("Remove remote stream {}", remote);
        self.sockets.write().remove(&remote);
    }

    pub fn is_listening(&self) -> bool {
        self.listening.load(Ordering::Acquire)
    }

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

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