use std::sync::Arc;
use tokio::process::{Child, ChildStdin, ChildStdout};
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use jsonrpsee::core::async_trait;
use jsonrpsee::core::client::{ReceivedMessage, TransportReceiverT, TransportSenderT};
use tracing::debug;

#[derive(Debug, Clone)]
pub struct StdioTransport {
    stdin: Arc<tokio::sync::Mutex<ChildStdin>>,
    stdout: Arc<tokio::sync::Mutex<BufReader<ChildStdout>>>,
}

impl StdioTransport {
    pub fn new(mut child: Child) -> Self {
        let stdin = Arc::new(tokio::sync::Mutex::new(child.stdin.take().unwrap()));
        let stdout = Arc::new(tokio::sync::Mutex::new(BufReader::new(child.stdout.take().unwrap())));
        Self { stdin, stdout }
    }
}

#[async_trait]
impl TransportSenderT for StdioTransport {
    type Error = tokio::io::Error;

    #[tracing::instrument(skip(self), level = "trace")]
    async fn send(&mut self, msg: String) -> Result<(), Self::Error> {
        debug!("Sending: {}", msg);
        let mut stdin = self.stdin.lock().await;
        stdin.write_all(msg.as_bytes()).await?;
        stdin.write_all(b"\n").await?;
        Ok(())
    }
}

#[async_trait]
impl TransportReceiverT for StdioTransport {
    type Error = tokio::io::Error;

    #[tracing::instrument(skip(self), level = "trace")]
    async fn receive(&mut self) -> Result<ReceivedMessage, Self::Error> {
        let mut stdout = self.stdout.lock().await;
        let mut str = String::new();
        stdout.read_line(&mut str).await?;
        debug!("Received: {}", str);
        Ok(ReceivedMessage::Text(str))
    }
}