//! Multi-line JSON reader for stdio-based transports.
//!
//! Provides resilient JSON-RPC message reading that handles agents or clients
//! sending multi-line (pretty-printed) JSON instead of strict single-line
//! JSON-RPC. Uses `serde_json::StreamDeserializer` to correctly extract
//! complete JSON values from a buffer, even when they span multiple lines
//! or when multiple values are concatenated.
//!
//! # Usage
//!
//! Both sides of a JSON-RPC stdio connection can use this:
//! - **Client side** (`connectors::acp::transport::StdioTransport`): reading agent stdout
//! - **Server side** (`dirigate`): reading client stdin
//!
//! ```ignore
//! use dirigent_core::acp::transport::json_reader::{JsonLineReader, ReadResult};
//! use tokio::io::BufReader;
//!
//! let stdin = tokio::io::stdin();
//! let mut reader = BufReader::new(stdin);
//! let mut json_reader = JsonLineReader::new();
//!
//! loop {
//!     match json_reader.read_message(&mut reader).await {
//!         Ok(ReadResult::Message(msg)) => println!("Got: {}", msg),
//!         Ok(ReadResult::Eof) => break,
//!         Err(e) => eprintln!("Error: {}", e),
//!     }
//! }
//! ```

use serde_json::Value;
use tokio::io::{AsyncBufRead, AsyncBufReadExt};

/// Safety limit: max lines to buffer before giving up.
const MAX_BUFFER_LINES: usize = 50_000;
/// Safety limit: max buffer size in bytes (10 MB).
const MAX_BUFFER_BYTES: usize = 10 * 1024 * 1024;

/// Result of reading one JSON message from the stream.
pub enum ReadResult {
    /// Successfully read a complete JSON message.
    Message(Value),
    /// EOF reached (stream closed).
    Eof,
}

/// Persistent state for a multi-line JSON reader.
///
/// Maintains a pending message buffer across `read_message` calls to handle
/// cases where the stream deserializer extracts two messages from one buffer.
pub struct JsonLineReader {
    /// Pending message from a previous parse that found two values in one buffer.
    pending: Option<Value>,
    /// Incomplete remainder from a previous parse where a complete JSON message
    /// was followed by the start of another message that wasn't yet complete.
    pending_remainder: Option<String>,
}

impl JsonLineReader {
    pub fn new() -> Self {
        Self {
            pending: None,
            pending_remainder: None,
        }
    }

    /// Read the next complete JSON message from the given buffered reader.
    ///
    /// Handles:
    /// - Single-line JSON (normal case)
    /// - Multi-line / pretty-printed JSON
    /// - Two concatenated JSON objects in the buffer
    /// - Non-JSON garbage (returns error immediately)
    ///
    /// The reader is borrowed mutably for each call but state is maintained
    /// in `self` across calls (for pending messages).
    pub async fn read_message<R: AsyncBufRead + Unpin>(
        &mut self,
        reader: &mut R,
    ) -> Result<ReadResult, String> {
        // Return pending message from previous parse if available
        if let Some(msg) = self.pending.take() {
            tracing::debug!("Returning pending message from previous read_message() parse");
            return Ok(ReadResult::Message(msg));
        }

        let mut json_buffer = String::new();
        let mut buffered_lines: usize = 0;

        // Initialize buffer from pending remainder if available.
        if let Some(remainder) = self.pending_remainder.take() {
            tracing::debug!(
                remainder_bytes = remainder.len(),
                "Initializing buffer from pending remainder"
            );
            json_buffer = remainder;
            buffered_lines = 1;
        }

        loop {
            let mut line = String::new();
            let bytes_read = reader
                .read_line(&mut line)
                .await
                .map_err(|e| format!("Failed to read line: {}", e))?;

            // EOF
            if bytes_read == 0 {
                if !json_buffer.is_empty() {
                    tracing::warn!(
                        buffer_lines = buffered_lines,
                        buffer_bytes = json_buffer.len(),
                        "EOF with incomplete multi-line JSON buffer"
                    );
                }
                return Ok(ReadResult::Eof);
            }

            // Skip empty lines when not buffering
            if line.trim().is_empty() && json_buffer.is_empty() {
                continue;
            }

            // Append to buffer
            if json_buffer.is_empty() {
                json_buffer = line;
                buffered_lines = 1;
            } else {
                json_buffer.push_str(&line);
                buffered_lines += 1;
            }

            // Try to extract a complete JSON value
            if let Some((message, remainder)) = try_extract_json(&json_buffer) {
                if buffered_lines > 1 {
                    tracing::warn!(
                        lines = buffered_lines,
                        bytes = json_buffer.len(),
                        "Recovered multi-line JSON message"
                    );
                }

                // If there's leftover data, try to parse as a second message
                if !remainder.is_empty() {
                    tracing::debug!(
                        remainder_bytes = remainder.len(),
                        "Buffer has remaining data after complete JSON"
                    );
                    if let Some((second_msg, second_remainder)) = try_extract_json(&remainder) {
                        self.pending = Some(second_msg);
                        if !second_remainder.is_empty() {
                            tracing::debug!(
                                remainder_bytes = second_remainder.len(),
                                "Carrying over third message fragment as pending remainder"
                            );
                            self.pending_remainder = Some(second_remainder);
                        }
                    } else {
                        // Incomplete second message — carry over instead of discarding
                        tracing::debug!(
                            remainder_bytes = remainder.len(),
                            "Carrying over incomplete remainder for next read_message()"
                        );
                        self.pending_remainder = Some(remainder);
                    }
                }

                return Ok(ReadResult::Message(message));
            }

            // Buffer doesn't contain a complete JSON value yet.
            // Check if it could ever become valid JSON.
            let buffer_trimmed = json_buffer.trim_start();
            if !buffer_trimmed.starts_with('{') && !buffer_trimmed.starts_with('[') {
                return Err(format!(
                    "Failed to parse JSON: expected JSON object. Line: {}",
                    &buffer_trimmed[..buffer_trimmed.len().min(200)]
                ));
            }

            // Safety limits
            if buffered_lines >= MAX_BUFFER_LINES {
                return Err(format!(
                    "Multi-line JSON message exceeded {} lines without completing",
                    MAX_BUFFER_LINES
                ));
            }
            if json_buffer.len() > MAX_BUFFER_BYTES {
                return Err(format!(
                    "Multi-line JSON message exceeded {} bytes without completing",
                    MAX_BUFFER_BYTES
                ));
            }

            if buffered_lines % 1000 == 0 {
                tracing::debug!(
                    lines = buffered_lines,
                    bytes = json_buffer.len(),
                    "Still buffering multi-line JSON message..."
                );
            }
        }
    }
}

/// Try to extract a complete JSON value from the front of `buffer`.
///
/// Uses `serde_json::StreamDeserializer` to parse the first complete JSON value.
/// Returns `Some((value, remaining))` if a value was found, where `remaining`
/// is the unparsed tail of the buffer. Returns `None` if the buffer doesn't
/// contain a complete JSON value yet.
pub fn try_extract_json(buffer: &str) -> Option<(Value, String)> {
    let mut stream = serde_json::Deserializer::from_str(buffer).into_iter::<Value>();
    match stream.next() {
        Some(Ok(value)) => {
            let byte_offset = stream.byte_offset();
            let remaining = buffer[byte_offset..].trim_start().to_string();
            Some((value, remaining))
        }
        _ => None,
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_extract_single_json() {
        let input = r#"{"jsonrpc":"2.0","id":1,"result":{}}"#;
        let (value, remainder) = try_extract_json(input).unwrap();
        assert_eq!(value["id"], 1);
        assert!(remainder.is_empty());
    }

    #[test]
    fn test_extract_incomplete_returns_none() {
        assert!(try_extract_json(r#"{"id":1,"res"#).is_none());
    }

    #[test]
    fn test_extract_two_concatenated() {
        let input = r#"{"id":1,"result":{}}
{"id":2,"result":{}}"#;
        let (first, remainder) = try_extract_json(input).unwrap();
        assert_eq!(first["id"], 1);
        let (second, remainder2) = try_extract_json(&remainder).unwrap();
        assert_eq!(second["id"], 2);
        assert!(remainder2.is_empty());
    }

    #[test]
    fn test_extract_multiline_pretty_printed() {
        let input = "{\n  \"id\": 1,\n  \"result\": {}\n}";
        let (value, remainder) = try_extract_json(input).unwrap();
        assert_eq!(value["id"], 1);
        assert!(remainder.is_empty());
    }

    #[test]
    fn test_extract_garbage_returns_none() {
        assert!(try_extract_json("on-management\",\n\n div {").is_none());
    }

    #[test]
    fn test_extract_empty_returns_none() {
        assert!(try_extract_json("").is_none());
        assert!(try_extract_json("   ").is_none());
    }

    #[test]
    fn test_extract_complete_followed_by_partial() {
        let input = r#"{"id":1,"result":{}}
{"id":2,"res"#;
        let (first, remainder) = try_extract_json(input).unwrap();
        assert_eq!(first["id"], 1);
        assert!(try_extract_json(&remainder).is_none());
    }

    #[test]
    fn test_extract_with_trailing_whitespace() {
        let input = r#"{"id":1,"result":{}}   "#;
        let (value, remainder) = try_extract_json(input).unwrap();
        assert_eq!(value["id"], 1);
        assert!(remainder.is_empty());
    }

    #[test]
    fn test_extract_with_escaped_newlines() {
        let input = r#"{"content":"line1\nline2\nline3"}"#;
        let (value, remainder) = try_extract_json(input).unwrap();
        assert_eq!(value["content"], "line1\nline2\nline3");
        assert!(remainder.is_empty());
    }

    #[tokio::test]
    async fn test_reader_single_line() {
        let input = b"{\"id\":1,\"result\":{}}\n";
        let mut cursor = std::io::Cursor::new(input.to_vec());
        let mut reader = JsonLineReader::new();
        match reader.read_message(&mut cursor).await.unwrap() {
            ReadResult::Message(msg) => assert_eq!(msg["id"], 1),
            ReadResult::Eof => panic!("Expected message, got EOF"),
        }
    }

    #[tokio::test]
    async fn test_reader_multiline() {
        let input = b"{\n  \"id\": 1,\n  \"result\": {}\n}\n";
        let mut cursor = std::io::Cursor::new(input.to_vec());
        let mut reader = JsonLineReader::new();
        match reader.read_message(&mut cursor).await.unwrap() {
            ReadResult::Message(msg) => assert_eq!(msg["id"], 1),
            ReadResult::Eof => panic!("Expected message, got EOF"),
        }
    }

    #[tokio::test]
    async fn test_reader_eof() {
        let input = b"";
        let mut cursor = std::io::Cursor::new(input.to_vec());
        let mut reader = JsonLineReader::new();
        match reader.read_message(&mut cursor).await.unwrap() {
            ReadResult::Message(_) => panic!("Expected EOF"),
            ReadResult::Eof => {}
        }
    }

    #[tokio::test]
    async fn test_reader_garbage_returns_error() {
        let input = b"not json at all\n";
        let mut cursor = std::io::Cursor::new(input.to_vec());
        let mut reader = JsonLineReader::new();
        let result = reader.read_message(&mut cursor).await;
        assert!(result.is_err());
        // Error should contain the raw line for forensics
        let err = result.err().unwrap();
        assert!(err.contains("not json at all"), "Error should contain raw line: {}", err);
    }

    #[tokio::test]
    async fn test_reader_carries_over_incomplete_remainder() {
        // Line 1 has a complete JSON object concatenated with the start of a second
        // (split at a JSON whitespace boundary — between tokens, not inside a string).
        // read_line returns `{"id":1}{"id":2,\n` (up to first newline).
        // try_extract_json extracts id:1, remainder `{"id":2,` is incomplete.
        // With the fix, this remainder is carried over as pending_remainder.
        // Second read_message() initializes json_buffer from it, reads the next line
        // `"result":{}}\n`, assembles `{"id":2,\n"result":{}}\n`, and extracts id:2.
        let input = b"{\"id\":1}{\"id\":2,\n\"result\":{}}\n";
        let mut cursor = std::io::Cursor::new(input.to_vec());
        let mut reader = JsonLineReader::new();

        // First call: extracts id:1, carries over incomplete remainder
        match reader.read_message(&mut cursor).await.unwrap() {
            ReadResult::Message(msg) => assert_eq!(msg["id"], 1),
            ReadResult::Eof => panic!("Expected message 1, got EOF"),
        }

        // Second call: assembles id:2 from remainder + next line
        match reader.read_message(&mut cursor).await.unwrap() {
            ReadResult::Message(msg) => assert_eq!(msg["id"], 2),
            ReadResult::Eof => panic!("Expected message 2, got EOF"),
        }
    }

    #[tokio::test]
    async fn test_reader_remainder_multiline_recovery() {
        // Simulates the bug scenario: agent sends streaming notifications rapidly,
        // OS pipe buffering delivers a complete message concatenated with the start
        // of the next on the same line. The second message spans multiple lines.
        //
        // Line 1: complete notification + start of second notification (no newline between)
        // Line 2: continuation of second notification's JSON
        // Line 3: closing of second notification
        let msg1 = r#"{"jsonrpc":"2.0","method":"notify","params":{"ok":true}}"#;
        // Second message is split across lines at JSON-whitespace boundaries
        let input = format!(
            "{}{{\n\"jsonrpc\":\"2.0\",\n\"method\":\"update\",\"params\":{{\"content\":\"card data\"}}}}\n",
            msg1
        );
        let mut cursor = std::io::Cursor::new(input.into_bytes());
        let mut reader = JsonLineReader::new();

        // First message
        match reader.read_message(&mut cursor).await.unwrap() {
            ReadResult::Message(msg) => {
                assert_eq!(msg["method"], "notify");
                assert_eq!(msg["params"]["ok"], true);
            }
            ReadResult::Eof => panic!("Expected message 1"),
        }

        // Second message (assembled from remainder + subsequent lines)
        match reader.read_message(&mut cursor).await.unwrap() {
            ReadResult::Message(msg) => {
                assert_eq!(msg["method"], "update");
                assert_eq!(msg["params"]["content"], "card data");
            }
            ReadResult::Eof => panic!("Expected message 2"),
        }
    }
}
