use anyhow::Result;
use async_trait::async_trait;
use futures_util::{SinkExt, StreamExt};
use tokio_tungstenite::tungstenite::Message;

use crate::transport::WsStream;

// ============================================================================
// Transport trait abstraction
// ============================================================================

/// Outbound half of a VPN transport connection.
#[async_trait]
pub trait TransportSink: Send + 'static {
    /// Send a framed binary message on the reliable channel.
    async fn send_reliable(&mut self, data: Vec<u8>) -> Result<()>;

    /// Send a datagram (unreliable, best-effort).
    /// Falls back to reliable if the transport does not support datagrams.
    async fn send_datagram(&mut self, data: Vec<u8>) -> Result<()>;

    /// Gracefully close the transport.
    async fn close(&mut self) -> Result<()>;
}

/// Inbound half of a VPN transport connection.
#[async_trait]
pub trait TransportStream: Send + 'static {
    /// Receive the next reliable binary message. Returns `None` on close.
    async fn recv_reliable(&mut self) -> Result<Option<Vec<u8>>>;

    /// Receive the next datagram. Returns `None` if datagrams are unsupported
    /// or the connection is closed.
    async fn recv_datagram(&mut self) -> Result<Option<Vec<u8>>>;

    /// Whether this transport supports unreliable datagrams.
    fn supports_datagrams(&self) -> bool;
}

// ============================================================================
// WebSocket implementation
// ============================================================================

/// WebSocket transport sink (wraps the write half of a split WsStream).
pub struct WsTransportSink {
    inner: futures_util::stream::SplitSink<WsStream, Message>,
}

impl WsTransportSink {
    pub fn new(inner: futures_util::stream::SplitSink<WsStream, Message>) -> Self {
        Self { inner }
    }
}

#[async_trait]
impl TransportSink for WsTransportSink {
    async fn send_reliable(&mut self, data: Vec<u8>) -> Result<()> {
        self.inner.send(Message::Binary(data.into())).await?;
        Ok(())
    }

    async fn send_datagram(&mut self, data: Vec<u8>) -> Result<()> {
        // WebSocket has no datagram support — fall back to reliable.
        self.send_reliable(data).await
    }

    async fn close(&mut self) -> Result<()> {
        self.inner.close().await?;
        Ok(())
    }
}

/// WebSocket transport stream (wraps the read half of a split WsStream).
pub struct WsTransportStream {
    inner: futures_util::stream::SplitStream<WsStream>,
}

impl WsTransportStream {
    pub fn new(inner: futures_util::stream::SplitStream<WsStream>) -> Self {
        Self { inner }
    }
}

#[async_trait]
impl TransportStream for WsTransportStream {
    async fn recv_reliable(&mut self) -> Result<Option<Vec<u8>>> {
        loop {
            match self.inner.next().await {
                Some(Ok(Message::Binary(data))) => return Ok(Some(data.to_vec())),
                Some(Ok(Message::Close(_))) | None => return Ok(None),
                Some(Ok(Message::Ping(_))) => {
                    // Ping handling is done at the tungstenite layer automatically
                    // when the sink side is alive. Just skip here.
                    continue;
                }
                Some(Ok(_)) => continue,
                Some(Err(e)) => return Err(anyhow::anyhow!("WebSocket error: {}", e)),
            }
        }
    }

    async fn recv_datagram(&mut self) -> Result<Option<Vec<u8>>> {
        // WebSocket does not support datagrams.
        Ok(None)
    }

    fn supports_datagrams(&self) -> bool {
        false
    }
}

/// Split a WebSocket stream into transport sink and stream halves.
pub fn split_ws(ws: WsStream) -> (WsTransportSink, WsTransportStream) {
    let (sink, stream) = ws.split();
    (WsTransportSink::new(sink), WsTransportStream::new(stream))
}