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

use alloc::{
    boxed::Box,
    collections::btree_map::BTreeMap,
    sync::{Arc, Weak},
};
use async_trait::async_trait;
use libk::{
    error::Error,
    task::runtime::maybe_timeout,
    vfs::{File, FileReadiness, Socket},
};
use libk_mm::PageBox;
use libk_util::{
    hash_table::DefaultHashTable,
    queue::BoundedMpmcQueue,
    sync::spin_rwlock::{IrqSafeRwLock, IrqSafeRwLockReadGuard},
};
use yggdrasil_abi::net::{types::LocalSocketAddress, MessageHeader, MessageHeaderMut};

use super::OwnedAddress;

struct Message {
    source: OwnedAddress,
    data: PageBox<[u8]>,
    ancillary: Option<Arc<File>>,
}

struct Pairing {
    address: OwnedAddress,
    socket: Weak<LocalPacketSocket>,
}

pub struct LocalPacketSocket {
    local: IrqSafeRwLock<Option<Arc<str>>>,
    connected: IrqSafeRwLock<Option<Arc<str>>>,
    remotes: IrqSafeRwLock<DefaultHashTable<OwnedAddress, Pairing>>,

    receive_queue: BoundedMpmcQueue<Message>,
    anonymous_socket_id: AtomicU64,
}

impl LocalPacketSocket {
    pub fn new() -> Arc<Self> {
        Arc::new(Self {
            local: IrqSafeRwLock::new(None),
            connected: IrqSafeRwLock::new(None),
            remotes: IrqSafeRwLock::new(DefaultHashTable::new()),

            receive_queue: BoundedMpmcQueue::new(128),
            anonymous_socket_id: AtomicU64::new(0),
        })
    }

    #[inline]
    fn make_message(message: &MessageHeader, source: OwnedAddress) -> Result<Message, Error> {
        let mut data = PageBox::new_slice(0, message.payload.len())?;
        let ancillary = super::read_ancillary(message)?;
        data.copy_from_slice(message.payload);
        Ok(Message {
            data,
            source,
            ancillary,
        })
    }

    fn push_nonblocking(&self, message: &MessageHeader, source: OwnedAddress) -> Result<(), Error> {
        let message = Self::make_message(message, source)?;
        self.receive_queue
            .try_push_back(message)
            .map_err(|_| Error::WouldBlock)?;
        Ok(())
    }

    async fn push(&self, message: &MessageHeader<'_>, source: OwnedAddress) -> Result<(), Error> {
        let message = Self::make_message(message, source)?;
        self.receive_queue.push_back(message).await;
        Ok(())
    }

    /// Makes the `socket` "known" to `self` by giving it an anonymous address.
    ///
    /// `self` must be a named socket (i.e., anonymous sockets cannot communicate to anonymous
    /// sockets).
    fn add_socket(&self, socket: &Arc<LocalPacketSocket>) -> Result<OwnedAddress, Error> {
        let path = self.local.read().clone().ok_or(Error::InvalidOperation)?;
        let mut remotes = self.remotes.write();
        let index = self.anonymous_socket_id.fetch_add(1, Ordering::Acquire);
        let address = OwnedAddress::Anonymous(index);
        assert!(!remotes.contains_key(&address));
        // Messages will be sent to the `socket` using the named address
        remotes.insert(
            address.clone(),
            Pairing {
                address: OwnedAddress::Path(path),
                socket: Arc::downgrade(socket),
            },
        );
        Ok(address)
    }

    /// Returns an address to use for sending data to the specified `remote` socket, and the
    /// socket itself.
    ///
    /// 1. If this socket is bound to a name, the name will be used
    /// 2. Otherwise:
    ///    2.1. If `remote` is "known" by this socket, the address in the table will be used
    ///    2.2. Otherwise, this socket is "paired" with the remote by generating a new anonymous
    ///    address unique to that socket
    fn get_send_info(
        self: &Arc<Self>,
        remote: &OwnedAddress,
    ) -> Result<(OwnedAddress, Arc<LocalPacketSocket>), Error> {
        // Find the socket first
        let peer = match remote {
            OwnedAddress::Path(path) => PATH_TABLE
                .read()
                .get(path)
                .cloned()
                .ok_or(Error::HostUnreachable)?,
            &OwnedAddress::Anonymous(anon) => self
                .remotes
                .read()
                .get(&OwnedAddress::Anonymous(anon))
                .and_then(|p| p.socket.upgrade())
                .ok_or(Error::HostUnreachable)?,
        };

        // (1)
        if let Some(address) = self.local.read().clone() {
            return Ok((OwnedAddress::Path(address), peer));
        }

        // (2.1)
        let remotes = self.remotes.read();
        if let Some(pairing) = remotes.get(remote) {
            return Ok((pairing.address.clone(), peer));
        }

        // (2.2)
        let mut remotes = IrqSafeRwLockReadGuard::upgrade(remotes);
        let address = peer.add_socket(self)?;
        // Insert pairing to the table so the socket "knows" how to send the messages
        // to the destination
        log::info!("local: now known to {remote:?} as {address:?}");
        remotes.insert(
            remote.clone(),
            Pairing {
                address: address.clone(),
                socket: Arc::downgrade(&peer),
            },
        );

        Ok((address, peer))
    }

    fn write_message(message: &mut MessageHeaderMut, msg: Message) -> Result<usize, Error> {
        let address = msg.source.to_borrowed();
        if message.payload.len() < msg.data.len() {
            log::warn!(
                "local: buffer ({}) too small to store a message of {}",
                message.payload.len(),
                msg.data.len()
            );
            return Err(Error::BufferTooSmall);
        }
        message.payload[..msg.data.len()].copy_from_slice(&msg.data);
        super::store_address(message, &address)?;
        super::write_ancillary(message, msg.ancillary)?;
        log::trace!("local: received {} bytes", msg.data.len());
        Ok(msg.data.len())
    }

    fn read_destination(&self, message: &MessageHeader) -> Result<OwnedAddress, Error> {
        if let Some(destination) = message.destination {
            super::load_address(destination).map(Into::into)
        } else if let Some(destination) = self.connected.read().clone() {
            Ok(OwnedAddress::Path(destination))
        } else {
            Err(Error::NotConnected)
        }
    }
}

static PATH_TABLE: IrqSafeRwLock<BTreeMap<Arc<str>, Arc<LocalPacketSocket>>> =
    IrqSafeRwLock::new(BTreeMap::new());

#[async_trait]
impl Socket for LocalPacketSocket {
    fn bind(self: Arc<Self>, local: &[u8]) -> Result<(), Error> {
        let local = super::load_address(local)?;

        let LocalSocketAddress::Path(local) = local else {
            // Cannot bind a local socket to an anonymous address
            return Err(Error::InvalidArgument);
        };

        let mut bound = self.local.write();
        if bound.is_some() {
            // Already bound
            return Err(Error::InvalidOperation);
        }

        let mut path_table = PATH_TABLE.write();
        if path_table.contains_key(local) {
            return Err(Error::AddrInUse);
        }

        log::info!("local: bound {local:?}");
        let local: Arc<str> = Arc::from(local);
        *bound = Some(local.clone());
        path_table.insert(local, self.clone());

        Ok(())
    }
    async fn connect(
        self: Arc<Self>,
        remote: &[u8],
        _timeout: Option<Duration>,
    ) -> Result<(), Error> {
        let remote = super::load_address(remote)?;

        let LocalSocketAddress::Path(remote) = remote else {
            // Cannot connect a local socket to an anonymous address
            return Err(Error::InvalidArgument);
        };

        let mut bound = self.connected.write();
        if bound.is_some() {
            return Err(Error::InvalidOperation);
        }

        let path: Arc<str> = remote.into();
        let address = OwnedAddress::Path(path.clone());
        // Associate with the address, discard the actual result
        let _ = self.get_send_info(&address)?;
        log::info!("local: connected to {remote:?}");
        *bound = Some(path);

        Ok(())
    }

    fn close(self: Arc<Self>) -> Result<(), Error> {
        let local = self.local.write().take();
        if let Some(path) = local {
            log::info!("local: unbound {path:?}");
            PATH_TABLE.write().remove(path.as_ref());
        }
        Ok(())
    }

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

    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> {
        let destination = self.read_destination(message)?;
        log::trace!(
            "local: send {} bytes to {:?}",
            message.payload.len(),
            destination
        );
        let (source, peer) = self.get_send_info(&destination)?;
        peer.push_nonblocking(message, source)?;
        Ok(message.payload.len())
    }
    async fn send_message(
        self: Arc<Self>,
        message: &MessageHeader,
        timeout: Option<Duration>,
    ) -> Result<usize, Error> {
        let _ = timeout;
        let destination = self.read_destination(message)?;
        log::trace!(
            "local: send {} bytes to {:?}",
            message.payload.len(),
            destination
        );
        let (source, peer) = self.get_send_info(&destination)?;
        peer.push(message, source).await?;
        Ok(message.payload.len())
    }

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

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

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