//! Session loop: read newline-delimited JSON-RPC envelopes, dispatch
//! each one, write the response back.
//!
//! Why generic over [`AsyncBufRead`] + [`AsyncWrite`]: in production the
//! CLI plugs in `tokio::io::stdin` and `tokio::io::stdout`. In tests we
//! plug in a `tokio::io::DuplexStream` pair so the entire IPC path is
//! exercised end to end without spawning a subprocess.
//!
//! Concurrency: ADR-0018 commits to single-threaded dispatch. Requests
//! land in arrival order, each one is awaited to completion, and the
//! next request is read off the inbound stream. The SDK can issue
//! multiple in-flight requests (id-correlated on its side) but the
//! co-process serializes them against the one MCP connection. This is
//! intentional: the upstream MCP server is one stateful peer, and the
//! v1.1 surface is a testing tool, not a high-throughput proxy.

use std::io;
use std::pin::Pin;

use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncWrite, AsyncWriteExt};

use super::dispatcher::{Dispatcher, McpBackend};
use super::protocol::{error_codes, IpcError, IpcRequest, IpcResponse};

/// Capacity ceiling the line buffer shrinks back to after an oversized
/// line (WOR-1313). 64 KiB comfortably holds every legitimate SDK
/// request while guaranteeing one hostile giant line cannot pin
/// `max_line_bytes` of heap for the rest of the session.
const LINE_BUFFER_RETAIN_BYTES: usize = 64 * 1024;

/// Tunable knobs for the session loop.
///
/// Kept as a struct so future fields (max-line length, idle timeout)
/// land additively. Defaults match the SDK contract: tolerate up to
/// 16 MiB envelopes (large `tools/list` responses) before bailing.
#[derive(Debug, Clone, Copy)]
pub struct SessionOptions {
    /// Maximum bytes allowed in a single line. A line longer than this
    /// is treated as a parse error and the connection closes. Set
    /// generously: the SDK never sends huge requests, but tools/list
    /// responses can carry a long catalog.
    pub max_line_bytes: usize,
}

impl Default for SessionOptions {
    fn default() -> Self {
        Self {
            max_line_bytes: 16 * 1024 * 1024,
        }
    }
}

/// Drive the IPC session until the inbound stream closes or the SDK
/// calls `mcp.shutdown`.
///
/// Returns once either condition is met. The caller is responsible for
/// any post-loop cleanup beyond what the backend's `shutdown()` already
/// did (for example flushing OS pipes).
pub async fn run<R, W, B>(
    inbound: R,
    mut outbound: W,
    dispatcher: Dispatcher<B>,
    options: SessionOptions,
) -> io::Result<()>
where
    R: AsyncBufRead + Unpin,
    W: AsyncWrite + Unpin,
    B: McpBackend,
{
    let mut inbound = inbound;
    let mut line = Vec::new();
    loop {
        match read_line_bounded(&mut inbound, &mut line, options.max_line_bytes).await? {
            LineRead::Eof => {
                // SDK closed stdin. Drain any reporter buffers and exit
                // cleanly. We ignore the backend's shutdown error because
                // there is no SDK left to surface it to.
                let _ = dispatcher.backend().shutdown().await;
                return Ok(());
            }
            LineRead::Oversized => {
                let envelope = IpcResponse::failure(
                    serde_json::Value::Null,
                    IpcError {
                        code: error_codes::PARSE_ERROR,
                        message: format!(
                            "envelope exceeded max_line_bytes ({})",
                            options.max_line_bytes
                        ),
                        data: None,
                    },
                );
                write_envelope(&mut outbound, &envelope).await?;
                // One giant line must not pin its capacity for the rest
                // of the session (WOR-1313).
                line.shrink_to(options.max_line_bytes.min(LINE_BUFFER_RETAIN_BYTES));
                continue;
            }
            LineRead::Line => {}
        }
        let trimmed = std::str::from_utf8(&line)
            .map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?
            .trim();
        if trimmed.is_empty() {
            continue;
        }
        let request: IpcRequest = match serde_json::from_str(trimmed) {
            Ok(r) => r,
            Err(err) => {
                let envelope = IpcResponse::failure(
                    serde_json::Value::Null,
                    IpcError {
                        code: error_codes::PARSE_ERROR,
                        message: format!("malformed json: {err}"),
                        data: None,
                    },
                );
                write_envelope(&mut outbound, &envelope).await?;
                continue;
            }
        };
        let is_shutdown = request.method == "mcp.shutdown";
        let response = dispatcher.dispatch(request).await;
        write_envelope(&mut outbound, &response).await?;
        if is_shutdown {
            return Ok(());
        }
    }
}

/// Outcome of one bounded line read.
enum LineRead {
    /// The stream ended with no pending bytes.
    Eof,
    /// `buf` holds one line (newline included when one was present).
    Line,
    /// The line exceeded the limit; its bytes were discarded, the
    /// stream is positioned after the terminating newline (or at EOF).
    Oversized,
}

/// Read one newline-terminated line into `buf`, refusing to buffer more
/// than `max_bytes` (WOR-1313).
///
/// Why not `read_line`: it appends the entire line before the caller
/// can check the size, so a hostile peer could grow the buffer to the
/// full line length (gigabytes) before the limit check ever ran. This
/// reader checks against the limit chunk by chunk and discards the
/// remainder of an oversized line while still consuming it from the
/// stream, keeping peak memory at `max_bytes` plus one BufReader chunk.
async fn read_line_bounded<R: AsyncBufRead + Unpin>(
    reader: &mut R,
    buf: &mut Vec<u8>,
    max_bytes: usize,
) -> io::Result<LineRead> {
    buf.clear();
    let mut oversized = false;
    loop {
        let (take, done) = {
            let chunk = reader.fill_buf().await?;
            if chunk.is_empty() {
                if buf.is_empty() && !oversized {
                    return Ok(LineRead::Eof);
                }
                return Ok(if oversized {
                    LineRead::Oversized
                } else {
                    LineRead::Line
                });
            }
            let (take, done) = match chunk.iter().position(|&b| b == b'\n') {
                Some(pos) => (pos + 1, true),
                None => (chunk.len(), false),
            };
            if !oversized {
                if buf.len() + take > max_bytes {
                    oversized = true;
                    buf.clear();
                } else {
                    buf.extend_from_slice(&chunk[..take]);
                }
            }
            (take, done)
        };
        Pin::new(&mut *reader).consume(take);
        if done {
            return Ok(if oversized {
                LineRead::Oversized
            } else {
                LineRead::Line
            });
        }
    }
}

async fn write_envelope<W: AsyncWrite + Unpin>(
    out: &mut W,
    response: &IpcResponse,
) -> io::Result<()> {
    let mut buf = serde_json::to_vec(response).map_err(io::Error::other)?;
    buf.push(b'\n');
    out.write_all(&buf).await?;
    out.flush().await?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::exec::protocol::IpcResponseBody;
    use async_trait::async_trait;
    use serde_json::json;
    use tokio::io::{duplex, BufReader};

    struct FlakyBackend;

    #[async_trait]
    impl McpBackend for FlakyBackend {
        async fn initialize(
            &self,
        ) -> Result<serde_json::Value, super::super::dispatcher::BackendError> {
            Ok(json!({"name": "flaky"}))
        }
        async fn call_tool(
            &self,
            _tool: &str,
            _arguments: serde_json::Value,
        ) -> Result<serde_json::Value, super::super::dispatcher::BackendError> {
            Err(super::super::dispatcher::BackendError::Server {
                message: "boom".into(),
                data: None,
            })
        }
        async fn list_tools(
            &self,
        ) -> Result<serde_json::Value, super::super::dispatcher::BackendError> {
            Ok(json!({"tools": []}))
        }
        async fn list_resources(
            &self,
        ) -> Result<serde_json::Value, super::super::dispatcher::BackendError> {
            Ok(json!({"resources": []}))
        }
        async fn read_resource(
            &self,
            _uri: &str,
        ) -> Result<serde_json::Value, super::super::dispatcher::BackendError> {
            Ok(json!({"contents": []}))
        }
        async fn call_prompt(
            &self,
            _name: &str,
            _arguments: serde_json::Value,
        ) -> Result<serde_json::Value, super::super::dispatcher::BackendError> {
            Ok(json!({"messages": []}))
        }
        async fn reporter_set_upload(
            &self,
            _url: &str,
            _token: Option<&str>,
        ) -> Result<(), super::super::dispatcher::BackendError> {
            Ok(())
        }
        async fn reporter_flush(&self) -> Result<(), super::super::dispatcher::BackendError> {
            Ok(())
        }
        async fn cache_set_mode(
            &self,
            _mode: &str,
        ) -> Result<(), super::super::dispatcher::BackendError> {
            Ok(())
        }
        async fn cassette_set_mode(
            &self,
            _mode: &str,
        ) -> Result<(), super::super::dispatcher::BackendError> {
            Ok(())
        }
        async fn shutdown(&self) -> Result<(), super::super::dispatcher::BackendError> {
            Ok(())
        }
    }

    #[tokio::test]
    async fn shutdown_ends_loop() {
        let (mut sdk, co_in) = duplex(8192);
        let (co_out, mut sdk_out) = duplex(8192);
        let dispatcher = Dispatcher::new(FlakyBackend);
        let handle = tokio::spawn(async move {
            run(
                BufReader::new(co_in),
                co_out,
                dispatcher,
                SessionOptions::default(),
            )
            .await
        });
        sdk.write_all(b"{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"mcp.shutdown\"}\n")
            .await
            .unwrap();
        let mut buf = vec![0u8; 1024];
        let n = sdk_out.read(&mut buf).await.unwrap();
        let line = std::str::from_utf8(&buf[..n]).unwrap();
        assert!(line.contains("\"result\""));
        handle.await.unwrap().unwrap();
    }

    #[tokio::test]
    async fn malformed_json_returns_parse_error_and_keeps_going() {
        let (mut sdk, co_in) = duplex(8192);
        let (co_out, mut sdk_out) = duplex(8192);
        let dispatcher = Dispatcher::new(FlakyBackend);
        let handle = tokio::spawn(async move {
            run(
                BufReader::new(co_in),
                co_out,
                dispatcher,
                SessionOptions::default(),
            )
            .await
        });
        sdk.write_all(b"not-json\n").await.unwrap();
        sdk.write_all(b"{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"mcp.initialize\"}\n")
            .await
            .unwrap();
        sdk.write_all(b"{\"jsonrpc\":\"2.0\",\"id\":3,\"method\":\"mcp.shutdown\"}\n")
            .await
            .unwrap();
        let mut all = Vec::new();
        loop {
            let mut chunk = vec![0u8; 1024];
            let n = sdk_out.read(&mut chunk).await.unwrap();
            if n == 0 {
                break;
            }
            all.extend_from_slice(&chunk[..n]);
            if all.iter().filter(|b| **b == b'\n').count() >= 3 {
                break;
            }
        }
        let text = String::from_utf8(all).unwrap();
        let lines: Vec<&str> = text.lines().collect();
        assert!(
            lines[0].contains("Parse")
                || lines[0].contains("parse")
                || lines[0].contains(&error_codes::PARSE_ERROR.to_string())
        );
        assert!(lines[1].contains("ipc_version"));
        let _ = handle.await;
    }

    use tokio::io::AsyncReadExt;

    #[tokio::test]
    async fn oversized_line_returns_parse_error_and_keeps_serving() {
        let (mut sdk, co_in) = duplex(8192);
        let (co_out, mut sdk_out) = duplex(8192);
        let dispatcher = Dispatcher::new(FlakyBackend);
        let options = SessionOptions {
            max_line_bytes: 128,
        };
        let handle =
            tokio::spawn(
                async move { run(BufReader::new(co_in), co_out, dispatcher, options).await },
            );
        // 1 KiB of garbage, way past the 128-byte cap, then a valid
        // shutdown request that must still be served.
        let mut giant = vec![b'x'; 1024];
        giant.push(b'\n');
        sdk.write_all(&giant).await.unwrap();
        sdk.write_all(b"{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"mcp.shutdown\"}\n")
            .await
            .unwrap();
        let mut all = Vec::new();
        loop {
            let mut chunk = vec![0u8; 1024];
            let n = sdk_out.read(&mut chunk).await.unwrap();
            if n == 0 {
                break;
            }
            all.extend_from_slice(&chunk[..n]);
            if all.iter().filter(|b| **b == b'\n').count() >= 2 {
                break;
            }
        }
        let text = String::from_utf8(all).unwrap();
        let lines: Vec<&str> = text.lines().collect();
        assert!(
            lines[0].contains("max_line_bytes (128)"),
            "first reply should flag the oversized envelope: {}",
            lines[0]
        );
        assert!(
            lines[1].contains("\"result\""),
            "session must keep serving after an oversized line: {}",
            lines[1]
        );
        handle.await.unwrap().unwrap();
    }

    #[tokio::test]
    async fn oversized_line_does_not_pin_buffer_capacity() {
        // WOR-1313: the bounded reader must never buffer more than
        // max_bytes of an oversized line, so one hostile giant line
        // cannot pin capacity for the rest of the session.
        let max = 256usize;
        let (mut sdk, co_in) = duplex(64 * 1024);
        let mut reader = BufReader::new(co_in);
        let writer = tokio::spawn(async move {
            let mut giant = vec![b'y'; 16 * 1024];
            giant.push(b'\n');
            sdk.write_all(&giant).await.unwrap();
            sdk.write_all(b"ok\n").await.unwrap();
            sdk
        });
        let mut buf = Vec::new();
        let outcome = read_line_bounded(&mut reader, &mut buf, max).await.unwrap();
        assert!(matches!(outcome, LineRead::Oversized));
        assert!(
            buf.capacity() <= max * 2,
            "oversized line grew the buffer to {} bytes (max {max})",
            buf.capacity()
        );
        // The remainder of the oversized line was consumed: the next
        // read yields the following line intact.
        let outcome = read_line_bounded(&mut reader, &mut buf, max).await.unwrap();
        assert!(matches!(outcome, LineRead::Line));
        assert_eq!(buf, b"ok\n");
        let _sdk = writer.await.unwrap();
    }

    #[tokio::test]
    async fn bounded_reader_handles_eof_without_trailing_newline() {
        let (mut sdk, co_in) = duplex(1024);
        sdk.write_all(b"tail-without-newline").await.unwrap();
        drop(sdk);
        let mut reader = BufReader::new(co_in);
        let mut buf = Vec::new();
        let outcome = read_line_bounded(&mut reader, &mut buf, 1024)
            .await
            .unwrap();
        assert!(matches!(outcome, LineRead::Line));
        assert_eq!(buf, b"tail-without-newline");
        let outcome = read_line_bounded(&mut reader, &mut buf, 1024)
            .await
            .unwrap();
        assert!(matches!(outcome, LineRead::Eof));
    }

    #[tokio::test]
    async fn dispatches_in_arrival_order() {
        let backend = FlakyBackend;
        let dispatcher = Dispatcher::new(backend);
        let r1 = dispatcher
            .dispatch(IpcRequest {
                jsonrpc: "2.0".into(),
                id: Some(json!(1)),
                method: "mcp.initialize".into(),
                params: None,
            })
            .await;
        let r2 = dispatcher
            .dispatch(IpcRequest {
                jsonrpc: "2.0".into(),
                id: Some(json!(2)),
                method: "mcp.call".into(),
                params: Some(json!({"tool": "x", "arguments": {}})),
            })
            .await;
        match r1.body {
            IpcResponseBody::Success { .. } => {}
            IpcResponseBody::Failure { error } => panic!("init failed: {error:?}"),
        }
        match r2.body {
            IpcResponseBody::Failure { error } => {
                assert_eq!(error.code, error_codes::MCP_SERVER_ERROR)
            }
            IpcResponseBody::Success { result } => panic!("expected failure, got {result:?}"),
        }
    }
}
