// SPDX-FileCopyrightText: 2024 blinry <mail@blinry.org>
// SPDX-FileCopyrightText: 2024 zormit <nt4u@kpvn.de>
// SPDX-FileCopyrightText: 2026 Caleb Maclennan <caleb@alerque.com>
//
// SPDX-License-Identifier: AGPL-3.0-or-later

//! Provides a way to write / read a socket through stdin, (un)packing content-length encoding.
//!
//! The idea is that a daemon process communicates through newline separated jsonrpc messages,
//! whereas LSP expects an HTTP-like Base Protocol:
//! <https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#baseProtocol>
//!
//! This forwarder thus
//! - takes jsonrpc from a socket (usually a daemon) and wraps it content-length encoded data to stdout
//! - takes content-length encoded data from stdin (as sent by an LSP client) and writes it
//!   "unpacked" to the socket

use std::path::Path;
use std::{env, process, str};

use anyhow::Context;
use async_trait::async_trait;
use futures::{SinkExt, StreamExt};
use tokio::io;
use tokio::io::{AsyncRead, AsyncWrite, BufReader, BufWriter};
use tokio::net::UnixStream;
use tokio::net::unix::{OwnedReadHalf, OwnedWriteHalf};
use tokio_util::bytes::{Buf, BytesMut};
use tokio_util::codec::{Decoder, Encoder, FramedRead, FramedWrite, LinesCodec};

use super::editor::strip_current_dir;

// TODO: Put these defaults to a module accessible by config.rs as well.
pub const DEFAULT_SOCKET_NAME: &str = "socket";
pub const CONFIG_DIR: &str = ".teamtype";

#[async_trait]
pub trait JSONRPCForwarder<
    R: AsyncRead + Unpin + Send + 'static,
    W: AsyncWrite + Unpin + Send + 'static,
>
{
    async fn connect_stream(
        &self,
        directory: &Path,
    ) -> anyhow::Result<(FramedRead<R, LinesCodec>, FramedWrite<W, LinesCodec>)>;

    async fn connection(&self, base_dir: &Path) -> anyhow::Result<()> {
        let (mut socket_read, mut socket_write) = self.connect_stream(base_dir).await?;

        // Construct stdin/stdout objects, which send/receive messages with a Content-Length header.
        let mut stdin = FramedRead::new(BufReader::new(io::stdin()), ContentLengthCodec);
        let mut stdout = FramedWrite::new(BufWriter::new(io::stdout()), ContentLengthCodec);

        tokio::spawn(async move {
            while let Some(Ok(message)) = socket_read.next().await {
                stdout
                    .send(message)
                    .await
                    .expect("Failed to write to stdout");
            }
            // Socket was closed.
            process::exit(0);
        });

        while let Some(Ok(message)) = stdin.next().await {
            socket_write.send(message).await?;
        }
        // Stdin was closed.
        process::exit(0);
    }
}

pub struct UnixJSONRPCForwarder {}

#[async_trait]
impl JSONRPCForwarder<OwnedReadHalf, OwnedWriteHalf> for UnixJSONRPCForwarder {
    async fn connect_stream(
        &self,
        directory: &Path,
    ) -> anyhow::Result<(
        FramedRead<OwnedReadHalf, LinesCodec>,
        FramedWrite<OwnedWriteHalf, LinesCodec>,
    )> {
        // Construct socket object, which send/receive newline-delimited messages.
        let socket_path = directory.join(CONFIG_DIR).join(DEFAULT_SOCKET_NAME);
        // See comment about SUN_LEN in editor.rs, but the TL;DR is that referencing a socket node
        // from a deeply nested or overly verbose path will fail on some platforms.
        // The extra song and dance to change into the parent directory first is not needed by our CLI
        // (which already changes to that location) but it will make this API usable when linked as a
        // library without changing the parent thread's location for keeps.
        let previous_cwd = env::current_dir()?;
        env::set_current_dir(
            socket_path
                .parent()
                .context("Invalid socket creation location")?,
        )?;
        let stream = UnixStream::connect(strip_current_dir(&socket_path)).await?;
        env::set_current_dir(previous_cwd)?;
        let (socket_read, socket_write) = stream.into_split();
        let socket_read = FramedRead::new(socket_read, LinesCodec::new());
        let socket_write = FramedWrite::new(socket_write, LinesCodec::new());
        Ok((socket_read, socket_write))
    }
}

struct ContentLengthCodec;

impl Encoder<String> for ContentLengthCodec {
    type Error = std::io::Error;

    fn encode(&mut self, item: String, dst: &mut BytesMut) -> Result<(), Self::Error> {
        let content_length = item.len();
        dst.extend_from_slice(format!("Content-Length: {content_length}\r\n\r\n").as_bytes());
        dst.extend_from_slice(item.as_bytes());
        Ok(())
    }
}

impl Decoder for ContentLengthCodec {
    type Error = anyhow::Error;
    type Item = String;

    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
        // Find the position of the Content-Length header.
        let c = b"Content-Length: ";
        let Some(start_of_header) = src.windows(c.len()).position(|window| window == c) else {
            return Ok(None);
        };

        // Find the end of the line after that.
        let (end_of_line, end_of_line_bytes) = match src[start_of_header + c.len()..]
            .windows(4)
            .position(|window| window == b"\r\n\r\n")
        {
            Some(pos) => (pos, 4),
            // Even though this is not valid in terms of the spec, also
            // accept plain newline separators in order to simplify manual testing.
            None => match src[start_of_header + c.len()..]
                .windows(2)
                .position(|window| window == b"\n\n")
            {
                Some(pos) => (pos, 2),
                None => return Ok(None),
            },
        };

        // Parse the content length.
        let content_length = str::from_utf8(
            &src[start_of_header + c.len()..start_of_header + c.len() + end_of_line],
        )?
        .parse()?;
        let content_start = start_of_header + c.len() + end_of_line + end_of_line_bytes;

        // Recommended optimization, in anticipation for future calls to `decode`.
        src.reserve(content_start + content_length);

        // Check if we have enough content.
        if src.len() < content_start + content_length {
            return Ok(None);
        }

        // Return the body of the message.
        src.advance(content_start);
        let content = src.split_to(content_length);
        Ok(Some(str::from_utf8(&content)?.to_string()))
    }
}
