use std::str;
use std::sync::Arc;

use anyhow::{Context as _, Result};
use collections::HashMap;
use futures::{
    AsyncBufReadExt, AsyncRead, AsyncReadExt as _, SinkExt as _,
    channel::mpsc::{Receiver, Sender, channel},
    io::BufReader,
};
use gpui::{BackgroundExecutor, Task};
use log::warn;
use parking_lot::Mutex;

use crate::{
    AnyResponse, CONTENT_LEN_HEADER, IoHandler, IoKind, NotificationOrRequest, RequestId,
    ResponseHandler,
};

const HEADER_DELIMITER: &[u8; 4] = b"\r\n\r\n";

/// Bounds the number of incoming LSP messages buffered between the background
/// reader and the foreground dispatcher. When the queue is full, the reader
/// stops reading the server's stdout, letting the OS pipe apply backpressure
/// to the server instead of buffering messages in memory without limit while
/// the foreground thread is unresponsive.
pub(crate) const INCOMING_MESSAGE_QUEUE_CAPACITY: usize = 128;

/// Handler for stdout of language server.
pub struct LspStdoutHandler {
    pub(super) loop_handle: Task<Result<()>>,
    pub(super) incoming_messages: Receiver<NotificationOrRequest>,
}

async fn read_headers<Stdout>(reader: &mut BufReader<Stdout>, buffer: &mut Vec<u8>) -> Result<()>
where
    Stdout: AsyncRead + Unpin + Send + 'static,
{
    loop {
        if buffer.len() >= HEADER_DELIMITER.len()
            && buffer[(buffer.len() - HEADER_DELIMITER.len())..] == HEADER_DELIMITER[..]
        {
            return Ok(());
        }

        if reader.read_until(b'\n', buffer).await? == 0 {
            anyhow::bail!("cannot read LSP message headers");
        }
    }
}

impl LspStdoutHandler {
    pub fn new<Input>(
        stdout: Input,
        response_handlers: Arc<Mutex<Option<HashMap<RequestId, ResponseHandler>>>>,
        io_handlers: Arc<Mutex<HashMap<i32, IoHandler>>>,
        cx: BackgroundExecutor,
    ) -> Self
    where
        Input: AsyncRead + Unpin + Send + 'static,
    {
        let (tx, notifications_channel) = channel(INCOMING_MESSAGE_QUEUE_CAPACITY);
        let loop_handle = cx.spawn(Self::handler(stdout, tx, response_handlers, io_handlers));
        Self {
            loop_handle,
            incoming_messages: notifications_channel,
        }
    }

    async fn handler<Input>(
        stdout: Input,
        mut notifications_sender: Sender<NotificationOrRequest>,
        response_handlers: Arc<Mutex<Option<HashMap<RequestId, ResponseHandler>>>>,
        io_handlers: Arc<Mutex<HashMap<i32, IoHandler>>>,
    ) -> anyhow::Result<()>
    where
        Input: AsyncRead + Unpin + Send + 'static,
    {
        let mut stdout = BufReader::new(stdout);

        let mut buffer = Vec::new();

        loop {
            buffer.clear();

            read_headers(&mut stdout, &mut buffer).await?;

            let headers = std::str::from_utf8(&buffer)?;

            let message_len = headers
                .split('\n')
                .find(|line| line.starts_with(CONTENT_LEN_HEADER))
                .and_then(|line| line.strip_prefix(CONTENT_LEN_HEADER))
                .with_context(|| format!("invalid LSP message header {headers:?}"))?
                .trim_end()
                .parse()?;

            buffer.resize(message_len, 0);
            stdout.read_exact(&mut buffer).await?;

            if let Ok(message) = str::from_utf8(&buffer) {
                log::trace!("incoming message: {message}");
                for handler in io_handlers.lock().values_mut() {
                    handler(IoKind::StdOut, message);
                }
            }

            if let Ok(msg) = serde_json::from_slice::<NotificationOrRequest>(&buffer) {
                notifications_sender.send(msg).await?;
            } else if let Ok(AnyResponse {
                id, error, result, ..
            }) = serde_json::from_slice(&buffer)
            {
                let handler = {
                    response_handlers
                        .lock()
                        .as_mut()
                        .and_then(|handlers| handlers.remove(&id))
                };
                if let Some(handler) = handler {
                    if let Some(error) = error {
                        handler(Err(error)).await;
                    } else if let Some(result) = result {
                        handler(Ok(result.get().into())).await;
                    } else {
                        handler(Ok("null".into())).await;
                    }
                }
            } else {
                warn!(
                    "failed to deserialize LSP message:\n{}",
                    std::str::from_utf8(&buffer)?
                );
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use futures::{AsyncWriteExt as _, StreamExt as _};
    use gpui::TestAppContext;

    #[gpui::test]
    async fn test_backpressure_when_messages_are_not_consumed(cx: &mut TestAppContext) {
        let total_messages = INCOMING_MESSAGE_QUEUE_CAPACITY * 4;
        let (mut writer, reader) = async_pipe::pipe();
        let mut handler = LspStdoutHandler::new(
            reader,
            Arc::new(Mutex::new(Some(HashMap::default()))),
            Arc::new(Mutex::new(HashMap::default())),
            cx.background_executor.clone(),
        );

        cx.background_executor
            .spawn(async move {
                let payload = r#"{"jsonrpc":"2.0","method":"test/notification","params":{}}"#;
                let message = format!("Content-Length: {}\r\n\r\n{}", payload.len(), payload);
                for _ in 0..total_messages {
                    writer.write_all(message.as_bytes()).await.unwrap();
                }
            })
            .detach();

        cx.run_until_parked();
        let mut received = 0;
        while handler.incoming_messages.try_recv().is_ok() {
            received += 1;
        }
        assert!(
            received < total_messages,
            "the reader buffered all {total_messages} messages while the consumer was wedged"
        );
        assert!(
            received <= INCOMING_MESSAGE_QUEUE_CAPACITY + 2,
            "expected at most {} buffered messages, got {received}",
            INCOMING_MESSAGE_QUEUE_CAPACITY + 2
        );

        while received < total_messages {
            assert!(
                handler.incoming_messages.next().await.is_some(),
                "the message stream ended after {received} of {total_messages} messages"
            );
            received += 1;
        }
    }

    #[gpui::test]
    async fn test_read_headers() {
        let mut buf = Vec::new();
        let mut reader = BufReader::new(b"Content-Length: 123\r\n\r\n" as &[u8]);
        read_headers(&mut reader, &mut buf).await.unwrap();
        assert_eq!(buf, b"Content-Length: 123\r\n\r\n");

        let mut buf = Vec::new();
        let mut reader = BufReader::new(b"Content-Type: application/vscode-jsonrpc\r\nContent-Length: 1235\r\n\r\n{\"somecontent\":123}" as &[u8]);
        read_headers(&mut reader, &mut buf).await.unwrap();
        assert_eq!(
            buf,
            b"Content-Type: application/vscode-jsonrpc\r\nContent-Length: 1235\r\n\r\n"
        );

        let mut buf = Vec::new();
        let mut reader = BufReader::new(b"Content-Length: 1235\r\nContent-Type: application/vscode-jsonrpc\r\n\r\n{\"somecontent\":true}" as &[u8]);
        read_headers(&mut reader, &mut buf).await.unwrap();
        assert_eq!(
            buf,
            b"Content-Length: 1235\r\nContent-Type: application/vscode-jsonrpc\r\n\r\n"
        );
    }
}