use crate::Result;
use serde_json::Value;
use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncReadExt, AsyncWrite, AsyncWriteExt};

#[derive(Debug, serde::Deserialize, thiserror::Error)]
#[error("{message}")]
pub(crate) struct RpcError {
    pub code: i64,
    pub message: String,
    #[serde(default)]
    pub data: Option<Value>,
}
impl RpcError {
    pub fn invalid(message: impl Into<String>) -> Self {
        Self {
            code: -32602,
            message: message.into(),
            data: None,
        }
    }
    pub fn failed(message: impl Into<String>) -> Self {
        Self {
            code: -32803,
            message: message.into(),
            data: None,
        }
    }
    pub fn modified() -> Self {
        Self {
            code: -32801,
            message: "Language-service item is stale; request it again".into(),
            data: None,
        }
    }
}

const MAX_HEADER: usize = 8192;
const MAX_BODY: usize = 64 * 1024 * 1024;

pub async fn read(reader: &mut (impl AsyncBufRead + Unpin)) -> Result<Option<Value>> {
    let mut header = Vec::new();
    loop {
        let available = reader.fill_buf().await?;
        if available.is_empty() {
            return if header.is_empty() {
                Ok(None)
            } else {
                Err("Truncated RPC header".into())
            };
        }
        let count = available
            .iter()
            .position(|b| *b == b'\n')
            .map_or(available.len(), |i| i + 1);
        if header.len() + count > MAX_HEADER {
            return Err("RPC header exceeds limit".into());
        }
        header.extend_from_slice(&available[..count]);
        reader.consume(count);
        if header.ends_with(b"\r\n\r\n") {
            break;
        }
    }
    let mut length = None;
    for line in std::str::from_utf8(&header)?
        .split("\r\n")
        .filter(|s| !s.is_empty())
    {
        let (name, value) = line.split_once(':').ok_or("Invalid RPC header")?;
        if name.eq_ignore_ascii_case("content-length") {
            if length.is_some() {
                return Err("Duplicate Content-Length".into());
            }
            length = Some(value.trim().parse::<usize>()?);
        }
    }
    let length = length.ok_or("Missing Content-Length")?;
    if length > MAX_BODY {
        return Err("RPC body exceeds limit".into());
    }
    let mut body = vec![0; length];
    reader.read_exact(&mut body).await?;
    let message: Value = serde_json::from_slice(&body)?;
    if !message.is_object() || message["jsonrpc"] != "2.0" {
        return Err("Invalid JSON-RPC message".into());
    }
    Ok(Some(message))
}

pub async fn write(writer: &mut (impl AsyncWrite + Unpin), mut message: Value) -> Result<()> {
    message["jsonrpc"] = Value::String("2.0".into());
    let body = serde_json::to_vec(&message)?;
    if body.len() > MAX_BODY {
        return Err("RPC body exceeds limit".into());
    }
    writer
        .write_all(format!("Content-Length: {}\r\n\r\n", body.len()).as_bytes())
        .await?;
    writer.write_all(&body).await?;
    writer.flush().await?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;
    use tokio::io::BufReader;
    #[tokio::test]
    async fn utf8_and_multiple_frames() {
        let mut bytes = Vec::new();
        write(&mut bytes, json!({"id":1,"result":"😀"}))
            .await
            .unwrap();
        write(&mut bytes, json!({"id":2,"result":"next"}))
            .await
            .unwrap();
        let mut input = BufReader::with_capacity(1, bytes.as_slice());
        assert_eq!(read(&mut input).await.unwrap().unwrap()["result"], "😀");
        assert_eq!(read(&mut input).await.unwrap().unwrap()["result"], "next");
        assert!(read(&mut input).await.unwrap().is_none());
    }
    #[tokio::test]
    async fn malformed_and_bounded() {
        for bytes in [
            b"Content-Length: 9\r\n\r\n{}".as_slice(),
            b"Content-Length: 1\r\nContent-Length: 1\r\n\r\n0",
            b"Content-Length: 999999999\r\n\r\n",
        ] {
            assert!(read(&mut BufReader::new(bytes)).await.is_err());
        }
        assert!(
            read(&mut BufReader::new(vec![b'x'; MAX_HEADER + 1].as_slice()))
                .await
                .is_err()
        );
    }
}