rzmq 0.5.20

use crate::error::ZmqError;
#[cfg(feature = "io-uring")]
use crate::io_uring_backend::connection_handler::OutgoingMessage;
#[cfg(feature = "io-uring")]
use crate::io_uring_backend::ops::UringOpRequest;
#[cfg(feature = "io-uring")]
use crate::io_uring_backend::ops::{WAKEUP_STATE_SIGNALED, WAKEUP_STATE_SLEEPING};
use crate::message::{FrameBatch, Msg};
use crate::socket::events::MonitorSender;
use crate::socket::options::SocketOptions;
use crate::socket::SocketEvent;
#[cfg(feature = "io-uring")]
use crate::uring;
use crate::Context;

use std::any::Any;
use std::fmt;
#[cfg(feature = "io-uring")]
use std::os::{fd::AsRawFd, unix::io::RawFd};
#[cfg(feature = "io-uring")]
use std::sync::atomic::{AtomicU8, Ordering};
use std::sync::Arc;
#[cfg(feature = "io-uring")]
use std::time::Duration;

use async_trait::async_trait;
use fibre::mpmc::AsyncSender;
#[cfg(feature = "io-uring")]
use fibre::mpsc;
#[cfg(feature = "io-uring")]
use fibre::oneshot::oneshot;
use fibre::{SendError, TrySendError};
use tokio::time::timeout as tokio_timeout;

#[async_trait]
pub(crate) trait ISocketConnection: Send + Sync + fmt::Debug {
  /// Sends a single message as a convenience wrapper around `send_multipart`.
  async fn send_message(&self, msg: Msg) -> Result<(), ZmqError> {
    let mut fb = FrameBatch::new();
    fb.push(msg);
    self.send_multipart(fb).await
  }

  /// Sends a complete logical message, which may consist of one or more parts.
  /// This is the primary method for sending data over the connection.
  async fn send_multipart(&self, msgs: FrameBatch) -> Result<(), ZmqError>;

  /// Attempts to send, returning ownership of `msgs` if the channel is immediately full
  /// (SNDTIMEO=0 path). For blocking/timed sends that time out, the message is consumed
  /// inside the dropped future and an empty batch is returned with the error.
  async fn send_multipart_owned(&self, msgs: FrameBatch) -> Result<(), (FrameBatch, ZmqError)> {
    match self.send_multipart(msgs).await {
      Ok(()) => Ok(()),
      Err(e) => Err((FrameBatch::new(), e)),
    }
  }

  /// Synchronous fast-path. Returns Err with ownership if full or closed.
  fn try_send_multipart_owned_sync(&self, msgs: FrameBatch) -> Result<(), (FrameBatch, ZmqError)> {
    Err((msgs, ZmqError::ResourceLimitReached)) // Default fallback
  }

  async fn close_connection(&self) -> Result<(), ZmqError>;
  fn as_any(&self) -> &dyn Any;
}

#[derive(Debug, Clone)]
pub(crate) struct DummyConnection;

#[async_trait]
impl ISocketConnection for DummyConnection {
  async fn send_message(&self, _msg: Msg) -> Result<(), ZmqError> {
    Err(ZmqError::UnsupportedFeature(
      "DummyConnection cannot send".into(),
    ))
  }

  async fn send_multipart(&self, _msgs: FrameBatch) -> Result<(), ZmqError> {
    Err(ZmqError::UnsupportedFeature(
      "DummyConnection cannot send multipart".into(),
    ))
  }

  async fn close_connection(&self) -> Result<(), ZmqError> {
    Ok(())
  }

  fn as_any(&self) -> &dyn Any {
    self
  }
}

#[cfg(feature = "io-uring")]
pub(crate) struct UringFdConnection {
  fd: RawFd,
  mpsc_tx: mpsc::BoundedAsyncSender<OutgoingMessage>,
  event_fd: eventfd::EventFD,
  worker_asleep: Arc<AtomicU8>,
  context: Context,
}

#[cfg(feature = "io-uring")]
impl fmt::Debug for UringFdConnection {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    f.debug_struct("UringFdConnection")
      .field("fd", &self.fd)
      .field("mpsc_tx_is_closed", &self.mpsc_tx.is_closed())
      .field("event_fd_raw", &self.event_fd.as_raw_fd())
      .field("context_present", &true)
      .finish()
  }
}

#[cfg(feature = "io-uring")]
impl UringFdConnection {
  pub(crate) fn new(
    fd: RawFd,
    mpsc_tx: mpsc::BoundedAsyncSender<OutgoingMessage>,
    event_fd: eventfd::EventFD,
    worker_asleep: Arc<AtomicU8>,
    context: Context,
  ) -> Self {
    Self {
      fd,
      mpsc_tx,
      event_fd,
      worker_asleep,
      context,
    }
  }
}

#[cfg(feature = "io-uring")]
impl UringFdConnection {
  async fn send_outgoing(&self, msg: OutgoingMessage) -> Result<(), ZmqError> {
    match self.mpsc_tx.send(msg).await {
      Ok(()) => {
        // Memory ordering contract (spec §4.1):
        //   Sender: mpsc_tx.send() [internal Release] → worker_asleep.load(Acquire)
        //   Worker: worker_asleep.store(true, Release) → re-drain MPSC → block
        // The Acquire load here pairs with the worker's Release store of `worker_asleep`.
        // Together they ensure: if the worker set `worker_asleep = true` before this load,
        // we will see it and write the eventfd. If we miss the store (false), the worker's
        // double-check re-drain loop will see our enqueued message and stay awake.
        if self.worker_asleep.load(Ordering::Relaxed) == WAKEUP_STATE_SLEEPING {
          if self.worker_asleep.compare_exchange(
            WAKEUP_STATE_SLEEPING,
            WAKEUP_STATE_SIGNALED,
            Ordering::AcqRel,
            Ordering::Acquire,
          ).is_ok() {
            if let Err(e) = self.event_fd.write(1) {
              tracing::error!("UringFdConnection: Failed to signal eventfd: {}", e);
            }
          }
        }
        Ok(())
      }
      Err(_) => Err(ZmqError::ConnectionClosed),
    }
  }
}

#[cfg(feature = "io-uring")]
#[async_trait]
impl ISocketConnection for UringFdConnection {
  
  async fn send_message(&self, msg: Msg) -> Result<(), ZmqError> {
    self.send_outgoing(OutgoingMessage::Single(msg)).await
  }

  async fn send_multipart(&self, msgs: FrameBatch) -> Result<(), ZmqError> {
    self.send_outgoing(OutgoingMessage::Multipart(msgs)).await
  }

  async fn close_connection(&self) -> Result<(), ZmqError> {
    // Closing is an async request to the worker.
    let (reply_tx, reply_rx) = oneshot();
    let unique_user_data = self.context.inner().next_handle() as u64;
    let req = UringOpRequest::ShutdownConnectionHandler {
      user_data: unique_user_data,
      fd: self.fd,
      reply_tx,
    };
    
    let worker_op_tx = uring::global_state::get_global_uring_worker_op_tx()?;
    worker_op_tx.send(req).await.map_err(|e| {
      ZmqError::Internal(format!("UringWorker op channel error for close: {}", e))
    })?;

    match tokio::time::timeout(Duration::from_secs(5), reply_rx.recv()).await {
      Ok(Ok(Ok(_))) => Ok(()),
      Ok(Ok(Err(e))) => Err(e),
      Ok(Err(_)) => Err(ZmqError::Internal("UringWorker reply channel error for close".into())),
      Err(_) => Err(ZmqError::Timeout),
    }
  }

  fn as_any(&self) -> &dyn Any {
    self
  }
}

#[derive(Clone)]
pub(crate) struct InprocConnection {
  connection_id: usize,
  local_pipe_write_id_to_peer: usize,
  local_pipe_read_id_from_peer: usize,
  peer_inproc_name_or_uri: String,
  context: Context,
  data_tx_to_peer: AsyncSender<FrameBatch>,
  monitor_tx: Option<MonitorSender>,
  socket_options: Arc<SocketOptions>,
}

impl fmt::Debug for InprocConnection {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    f.debug_struct("InprocConnection")
      .field("connection_id", &self.connection_id)
      .field(
        "local_pipe_write_id_to_peer",
        &self.local_pipe_write_id_to_peer,
      )
      .field(
        "local_pipe_read_id_from_peer",
        &self.local_pipe_read_id_from_peer,
      )
      .field("peer_inproc_name_or_uri", &self.peer_inproc_name_or_uri)
      .field("context_present", &true) // Context doesn't have a simple Debug
      .field("data_tx_to_peer_closed", &self.data_tx_to_peer.is_closed())
      .field("monitor_tx_is_some", &self.monitor_tx.is_some())
      .field("socket_options", &self.socket_options)
      .finish()
  }
}

impl InprocConnection {
  pub(crate) fn new(
    connection_id: usize,
    local_pipe_write_id_to_peer: usize,
    local_pipe_read_id_from_peer: usize,
    peer_inproc_name_or_uri: String,
    context: Context,
    data_tx_to_peer: AsyncSender<FrameBatch>,
    monitor_tx: Option<MonitorSender>,
    socket_options: Arc<SocketOptions>,
  ) -> Self {
    Self {
      connection_id,
      local_pipe_write_id_to_peer,
      local_pipe_read_id_from_peer,
      peer_inproc_name_or_uri,
      context,
      data_tx_to_peer,
      monitor_tx,
      socket_options,
    }
  }
}

#[async_trait]
impl ISocketConnection for InprocConnection {
  async fn send_multipart(&self, msgs: FrameBatch) -> Result<(), ZmqError> {
    let timeout_opt = self.socket_options.sndtimeo;

    match timeout_opt {
      None => {
        self.data_tx_to_peer.send(msgs).await.map_err(|_| {
          tracing::warn!(conn_id = self.connection_id, peer = %self.peer_inproc_name_or_uri, "InprocConnection send failed (ConnectionClosed)");
          ZmqError::ConnectionClosed
        })
      }
      Some(d) if d.is_zero() => {
        match self.data_tx_to_peer.try_send(msgs) {
          Ok(()) => Ok(()),
          Err(TrySendError::Full(_)) => Err(ZmqError::ResourceLimitReached),
          Err(TrySendError::Closed(_)) => {
            tracing::warn!(conn_id = self.connection_id, peer = %self.peer_inproc_name_or_uri, "InprocConnection non-blocking send failed (ConnectionClosed)");
            Err(ZmqError::ConnectionClosed)
          }
          _ => unreachable!(),
        }
      }
      Some(duration) => {
        match self.data_tx_to_peer.try_send(msgs) {
          Ok(()) => Ok(()),
          Err(TrySendError::Closed(_)) => {
            tracing::warn!(conn_id = self.connection_id, peer = %self.peer_inproc_name_or_uri, "InprocConnection send failed (ConnectionClosed)");
            Err(ZmqError::ConnectionClosed)
          }
          Err(TrySendError::Full(returned_msgs)) => {
            match tokio_timeout(duration, self.data_tx_to_peer.send(returned_msgs)).await {
              Ok(Ok(())) => Ok(()),
              Ok(Err(SendError::Closed)) => {
                tracing::warn!(conn_id = self.connection_id, peer = %self.peer_inproc_name_or_uri, "InprocConnection timed send failed (ConnectionClosed)");
                Err(ZmqError::ConnectionClosed)
              }
              Err(_) => Err(ZmqError::Timeout),
              _ => unreachable!(),
            }
          }
          _ => unreachable!(),
        }
      }
    }
  }

  async fn close_connection(&self) -> Result<(), ZmqError> {
    tracing::debug!(
      conn_id = self.connection_id,
      peer = %self.peer_inproc_name_or_uri,
      local_read_pipe_id_being_closed = self.local_pipe_read_id_from_peer,
      "InprocConnection::close_connection called."
    );

    if let Some(ref monitor) = self.monitor_tx {
      let event = SocketEvent::Disconnected {
        endpoint: self.peer_inproc_name_or_uri.clone(),
      };
      if monitor.try_send(event).is_err() {
        tracing::warn!(
          conn_id = self.connection_id,
          peer = %self.peer_inproc_name_or_uri,
          "Failed to send Disconnected monitor event for inproc connection (channel full/closed)."
        );
      } else {
        tracing::debug!(
          conn_id = self.connection_id,
          peer = %self.peer_inproc_name_or_uri,
          "Sent Disconnected monitor event for inproc connection."
        );
      }
    }

    // SCAX-backed inproc connections handle teardown via EOF on the DuplexStream.
    // No explicit peer-closed event needed.
    Ok(())
  }
  
  fn as_any(&self) -> &dyn Any {
    self
  }
}