// SPDX-License-Identifier: AGPL-3.0-or-later

use anyhow::Context;
use bitcode::{Decode, Encode, decode, encode};
use tokio::{
    io::{AsyncReadExt, AsyncWriteExt},
    net::{UnixListener, UnixStream},
};

#[derive(Debug, thiserror::Error)]
pub enum TransportError {
    #[error(transparent)]
    Context(#[from] anyhow::Error),
    #[error("io: {0}")]
    Io(#[from] std::io::Error),
    #[error("decode: {0}")]
    Decode(#[source] bitcode::Error),
}

pub trait Listener {
    type Transport: Transport;

    fn accept(&self) -> impl Future<Output = Result<Self::Transport, std::io::Error>> + Send;
}

impl Listener for UnixListener {
    type Transport = UnixStream;

    async fn accept(&self) -> Result<Self::Transport, std::io::Error> {
        self.accept().await.map(|(socket, _addr)| socket)
    }
}

pub trait Transport {
    fn send(
        &mut self,
        message: &(impl Encode + Send + Sync),
    ) -> impl Future<Output = Result<(), TransportError>> + Send;
    fn recv<T: for<'a> Decode<'a>>(
        &mut self,
    ) -> impl Future<Output = Result<T, TransportError>> + Send;
}

impl Transport for UnixStream {
    async fn send(&mut self, message: &(impl Encode + Send + Sync)) -> Result<(), TransportError> {
        self.writable().await.map_err(TransportError::Io)?;
        let encoded = encode(message);
        self.write_u64(encoded.len() as u64)
            .await
            .map_err(TransportError::Io)?;
        self.write_all(&encoded).await.map_err(TransportError::Io)?;
        Ok(())
    }

    async fn recv<T: for<'a> Decode<'a>>(&mut self) -> Result<T, TransportError> {
        self.readable()
            .await
            .context("when waiting for socket to be readable")?;
        let length = self
            .read_u64()
            .await
            .context("when reading message length")?;
        let mut buffer = vec![0; length as usize];
        self.read_exact(&mut buffer)
            .await
            .context("when reading message")?;

        decode(&buffer).map_err(TransportError::Decode)
    }
}
