simplify_baml 0.2.0

/// Simplified partial JSON parser for streaming support
///
/// Inspired by BAML's `jsonish` library but simplified for our use case.
/// Handles incomplete JSON objects and arrays from streaming LLM responses.

use anyhow::Result;
use serde_json::Value as JsonValue;

/// Normalize Unicode curly/smart quotes to ASCII quotes ONLY outside string literals
/// LLMs sometimes use curly quotes for JSON structure, which breaks parsing.
/// But curly quotes INSIDE strings should be preserved (escaped if needed).
fn normalize_quotes(text: &str) -> String {
    let mut result = String::with_capacity(text.len());
    let mut in_string = false;
    let mut escape_next = false;

    for c in text.chars() {
        if escape_next {
            result.push(c);
            escape_next = false;
            continue;
        }

        if c == '\\' && in_string {
            result.push(c);
            escape_next = true;
            continue;
        }

        if c == '"' {
            in_string = !in_string;
            result.push(c);
            continue;
        }

        if in_string {
            // Inside a string, escape curly quotes instead of converting them
            match c {
                '\u{201C}' | '\u{201D}' => result.push_str("\\\""),
                '\u{2018}' | '\u{2019}' => result.push('\''),
                '\n' => result.push_str("\\n"),
                '\r' => result.push_str("\\r"),
                '\t' => result.push_str("\\t"),
                _ => result.push(c),
            }
        } else {
            // Outside strings, convert curly quotes to ASCII for JSON structure
            match c {
                '\u{201C}' | '\u{201D}' => result.push('"'),
                '\u{2018}' | '\u{2019}' => result.push('\''),
                _ => result.push(c),
            }
        }
    }

    result
}

/// Attempt to parse potentially incomplete JSON from a streaming response
///
/// This function tries multiple strategies to handle partial JSON:
/// 1. Parse as-is (might already be valid JSON)
/// 2. Normalize curly/smart quotes (LLMs sometimes produce these)
/// 3. Escape unescaped control characters in strings
/// 4. Extract from markdown code blocks (only if above attempts fail)
/// 5. Auto-close open structures (for streaming partial responses)
///
/// NOTE: Markdown extraction is intentionally done AFTER trying to parse as raw JSON.
/// This prevents issues where valid JSON containing embedded markdown code blocks
/// (e.g., in a string field like `message`) would be incorrectly extracted.
///
/// # Arguments
/// * `partial_json` - The potentially incomplete JSON string
///
/// # Returns
/// * `Ok(Some(JsonValue))` - Successfully parsed (complete or auto-closed)
/// * `Ok(None)` - Cannot parse yet, need more data
/// * `Err(...)` - Unrecoverable parsing error
pub fn try_parse_partial_json(partial_json: &str) -> Result<Option<JsonValue>> {
    let trimmed = partial_json.trim();

    if trimmed.is_empty() {
        return Ok(None);
    }

    // Try parsing the raw input first (might already be valid JSON)
    // This handles the common case where LLM returns raw JSON without markdown
    if let Ok(value) = serde_json::from_str::<JsonValue>(trimmed) {
        return Ok(Some(value));
    }

    // Try with quote normalization (LLMs sometimes use curly/smart quotes)
    let normalized = normalize_quotes(trimmed);
    if let Ok(value) = serde_json::from_str::<JsonValue>(&normalized) {
        return Ok(Some(value));
    }

    // Try with control char escaping
    let escaped = escape_control_chars_in_strings(&normalized);
    if let Ok(value) = serde_json::from_str::<JsonValue>(&escaped) {
        return Ok(Some(value));
    }

    // Only now try markdown extraction as a fallback
    // This handles responses wrapped in ```json code blocks
    let extracted = extract_from_markdown(trimmed);
    if extracted != trimmed {
        // Markdown extraction changed something, try parsing the extracted content
        if let Ok(value) = serde_json::from_str::<JsonValue>(&extracted) {
            return Ok(Some(value));
        }

        // Try with normalization on extracted content
        let normalized_extracted = normalize_quotes(&extracted);
        if let Ok(value) = serde_json::from_str::<JsonValue>(&normalized_extracted) {
            return Ok(Some(value));
        }

        let escaped_extracted = escape_control_chars_in_strings(&normalized_extracted);
        if let Ok(value) = serde_json::from_str::<JsonValue>(&escaped_extracted) {
            return Ok(Some(value));
        }

        // Try auto-closing on the extracted content (for incomplete JSON in markdown)
        let attempts_extracted = generate_completion_attempts(&escaped_extracted);
        for attempt in attempts_extracted {
            if let Ok(value) = serde_json::from_str::<JsonValue>(&attempt) {
                return Ok(Some(value));
            }
        }
    }

    // Try auto-closing structures (for streaming partial responses)
    let attempts = generate_completion_attempts(&escaped);

    for attempt in attempts {
        if let Ok(value) = serde_json::from_str::<JsonValue>(&attempt) {
            return Ok(Some(value));
        }
    }

    // If we can't parse it yet, return None (need more data)
    Ok(None)
}

/// Extract JSON from markdown code blocks
fn extract_from_markdown(text: &str) -> String {
    // Check for ```json blocks (case-insensitive)
    let text_lower = text.to_lowercase();
    if let Some(start) = text_lower.find("```json") {
        let json_start = start + 7;
        if let Some(end_offset) = text[json_start..].find("```") {
            let json_end = json_start + end_offset;
            return text[json_start..json_end].trim().to_string();
        }
        // If no closing ```, return everything after ```json
        return text[json_start..].trim().to_string();
    }

    // Check for ``` blocks without language
    if let Some(start) = text.find("```") {
        let content_start = start + 3;
        if let Some(end) = text[content_start..].find("```") {
            let content_end = content_start + end;
            let content = text[content_start..content_end].trim();
            if content.starts_with('{') || content.starts_with('[') {
                return content.to_string();
            }
        } else {
            // No closing ```, return everything after ```
            let content = text[content_start..].trim();
            if content.starts_with('{') || content.starts_with('[') {
                return content.to_string();
            }
        }
    }

    // Try to find JSON boundaries
    if let Some(start) = text.find('{') {
        if let Some(end) = text.rfind('}') {
            if end > start {
                return text[start..=end].to_string();
            }
        }
        // No closing brace, return from { to end
        return text[start..].to_string();
    }

    if let Some(start) = text.find('[') {
        if let Some(end) = text.rfind(']') {
            if end > start {
                return text[start..=end].to_string();
            }
        }
        // No closing bracket, return from [ to end
        return text[start..].to_string();
    }

    text.to_string()
}

/// Escape literal control characters (newlines, tabs, etc.) inside JSON strings
///
/// LLMs often produce JSON with unescaped newlines/tabs inside string values.
/// This function walks through the JSON, tracking string boundaries, and escapes
/// any literal control characters found inside strings.
fn escape_control_chars_in_strings(json: &str) -> String {
    let mut result = String::with_capacity(json.len());
    let mut in_string = false;
    let mut escape_next = false;

    for c in json.chars() {
        if escape_next {
            result.push(c);
            escape_next = false;
            continue;
        }

        if c == '\\' && in_string {
            result.push(c);
            escape_next = true;
            continue;
        }

        if c == '"' {
            in_string = !in_string;
            result.push(c);
            continue;
        }

        if in_string {
            match c {
                '\n' => result.push_str("\\n"),
                '\r' => result.push_str("\\r"),
                '\t' => result.push_str("\\t"),
                '\x08' => result.push_str("\\b"),
                '\x0C' => result.push_str("\\f"),
                _ => result.push(c),
            }
        } else {
            result.push(c);
        }
    }

    result
}

/// Count braces/brackets outside of string literals
fn count_braces_string_aware(json: &str) -> (usize, usize, usize, usize) {
    let mut open_braces = 0;
    let mut close_braces = 0;
    let mut open_brackets = 0;
    let mut close_brackets = 0;
    let mut in_string = false;
    let mut escape_next = false;

    for c in json.chars() {
        if escape_next {
            escape_next = false;
            continue;
        }

        if c == '\\' && in_string {
            escape_next = true;
            continue;
        }

        if c == '"' {
            in_string = !in_string;
            continue;
        }

        if in_string {
            continue;
        }

        match c {
            '{' => open_braces += 1,
            '}' => close_braces += 1,
            '[' => open_brackets += 1,
            ']' => close_brackets += 1,
            _ => {}
        }
    }

    (open_braces, close_braces, open_brackets, close_brackets)
}

/// Generate multiple attempts to complete partial JSON
fn generate_completion_attempts(json: &str) -> Vec<String> {
    let json = json.trim();
    let mut attempts = Vec::new();

    // Strategy 1: Count braces/brackets (string-aware) and auto-close
    let (open_braces, close_braces, open_brackets, close_brackets) =
        count_braces_string_aware(json);

    let mut completion = json.to_string();

    // Check if we have an incomplete string at the end
    if has_incomplete_string(&completion) {
        completion.push('"');
    }

    // Close arrays first (inner to outer)
    for _ in 0..(open_brackets.saturating_sub(close_brackets)) {
        completion.push(']');
    }

    // Close objects
    for _ in 0..(open_braces.saturating_sub(close_braces)) {
        completion.push('}');
    }

    attempts.push(completion);

    // Strategy 2: More aggressive - assume we're in the middle of writing a value
    let mut aggressive = json.to_string();

    // If ends with a colon, comma, or opening bracket/brace, might be waiting for a value
    if json.trim_end().ends_with(':') {
        aggressive.push_str("null");
    } else if json.trim_end().ends_with(',') {
        // Remove trailing comma and try closing
        aggressive = aggressive.trim_end().trim_end_matches(',').to_string();
    }

    // Close incomplete string
    if has_incomplete_string(&aggressive) {
        aggressive.push('"');
    }

    // Close structures
    for _ in 0..(open_brackets.saturating_sub(close_brackets)) {
        aggressive.push(']');
    }
    for _ in 0..(open_braces.saturating_sub(close_braces)) {
        aggressive.push('}');
    }

    attempts.push(aggressive);

    // Strategy 3: Remove incomplete last field/element
    if let Some(last_comma) = json.rfind(',') {
        let mut truncated = json[..=last_comma].to_string();
        truncated = truncated.trim_end().trim_end_matches(',').to_string();

        for _ in 0..(open_brackets.saturating_sub(close_brackets)) {
            truncated.push(']');
        }
        for _ in 0..(open_braces.saturating_sub(close_braces)) {
            truncated.push('}');
        }

        attempts.push(truncated);
    }

    attempts
}

/// Check if the JSON string has an incomplete string value at the end
fn has_incomplete_string(json: &str) -> bool {
    let mut in_string = false;
    let mut escape_next = false;
    let mut last_quote_pos = None;

    for (i, c) in json.chars().enumerate() {
        if escape_next {
            escape_next = false;
            continue;
        }

        match c {
            '\\' if in_string => escape_next = true,
            '"' => {
                in_string = !in_string;
                if in_string {
                    last_quote_pos = Some(i);
                } else {
                    last_quote_pos = None;
                }
            }
            _ => {}
        }
    }

    // If we're in a string at the end, it's incomplete
    in_string && last_quote_pos.is_some()
}

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

    #[test]
    fn test_complete_json() {
        let json = r#"{"name": "John", "age": 30}"#;
        let result = try_parse_partial_json(json).unwrap();
        assert!(result.is_some());

        let value = result.unwrap();
        assert_eq!(value["name"], "John");
        assert_eq!(value["age"], 30);
    }

    #[test]
    fn test_incomplete_object() {
        let partial = r#"{"name": "John", "age": 30"#;
        let result = try_parse_partial_json(partial).unwrap();
        assert!(result.is_some());

        let value = result.unwrap();
        assert_eq!(value["name"], "John");
        assert_eq!(value["age"], 30);
    }

    #[test]
    fn test_incomplete_string() {
        let partial = r#"{"name": "Joh"#;
        let result = try_parse_partial_json(partial).unwrap();
        assert!(result.is_some());

        let value = result.unwrap();
        assert_eq!(value["name"], "Joh");
    }

    #[test]
    fn test_incomplete_array() {
        let partial = r#"{"items": [1, 2, 3"#;
        let result = try_parse_partial_json(partial).unwrap();
        assert!(result.is_some());

        let value = result.unwrap();
        assert_eq!(value["items"].as_array().unwrap().len(), 3);
    }

    #[test]
    fn test_nested_incomplete() {
        let partial = r#"{"person": {"name": "John", "age": 30"#;
        let result = try_parse_partial_json(partial).unwrap();
        assert!(result.is_some());

        let value = result.unwrap();
        assert_eq!(value["person"]["name"], "John");
        assert_eq!(value["person"]["age"], 30);
    }

    #[test]
    fn test_markdown_extraction() {
        let partial = r#"Here's the data:
```json
{"name": "John", "age": 30
```"#;
        let result = try_parse_partial_json(partial).unwrap();
        assert!(result.is_some());
    }

    #[test]
    fn test_markdown_incomplete() {
        let partial = r#"```json
{"name": "John", "age": 30"#;
        let result = try_parse_partial_json(partial).unwrap();
        assert!(result.is_some());
    }

    #[test]
    fn test_trailing_comma() {
        let partial = r#"{"name": "John", "age": 30,"#;
        let result = try_parse_partial_json(partial).unwrap();
        assert!(result.is_some());

        let value = result.unwrap();
        assert_eq!(value["name"], "John");
    }

    #[test]
    fn test_empty_input() {
        let result = try_parse_partial_json("").unwrap();
        assert!(result.is_none());
    }

    #[test]
    fn test_incomplete_field_name() {
        // This case is harder - we have an incomplete field being written
        let partial = r#"{"name": "John", "ag"#;
        let _result = try_parse_partial_json(partial).unwrap();
        // Should return Some with just the complete fields, or None if can't parse
        // Either is acceptable for this edge case
    }

    #[test]
    fn test_uppercase_json_code_fence() {
        let partial = r#"```JSON
{"name": "John", "age": 30}
```"#;
        let result = try_parse_partial_json(partial).unwrap();
        assert!(result.is_some());
        let value = result.unwrap();
        assert_eq!(value["name"], "John");
        assert_eq!(value["age"], 30);
    }

    #[test]
    fn test_mixed_case_json_code_fence() {
        let partial = r#"```Json
{"name": "Alice"}
```"#;
        let result = try_parse_partial_json(partial).unwrap();
        assert!(result.is_some());
        let value = result.unwrap();
        assert_eq!(value["name"], "Alice");
    }

    #[test]
    fn test_uppercase_json_code_fence_incomplete() {
        let partial = r#"```JSON
{"name": "John", "age": 30"#;
        let result = try_parse_partial_json(partial).unwrap();
        assert!(result.is_some());
        let value = result.unwrap();
        assert_eq!(value["name"], "John");
    }

    #[test]
    fn test_braces_inside_string_value() {
        let partial = r#"{"text": "use { for scope"#;
        let result = try_parse_partial_json(partial).unwrap();
        assert!(result.is_some());
        let value = result.unwrap();
        assert_eq!(value["text"], "use { for scope");
    }

    #[test]
    fn test_brackets_inside_string_value() {
        let partial = r#"{"code": "arr = [1, 2, 3]"#;
        let result = try_parse_partial_json(partial).unwrap();
        assert!(result.is_some());
        let value = result.unwrap();
        assert_eq!(value["code"], "arr = [1, 2, 3]");
    }

    #[test]
    fn test_mixed_braces_brackets_in_string() {
        let partial = r#"{"message": "JSON: {\"arr\": [1, 2]}"#;
        let result = try_parse_partial_json(partial).unwrap();
        assert!(result.is_some());
        let value = result.unwrap();
        assert_eq!(value["message"], r#"JSON: {"arr": [1, 2]}"#);
    }

    #[test]
    fn test_complete_json_with_braces_in_string() {
        let json = r#"{"text": "use { for scope and } to close", "count": 5}"#;
        let result = try_parse_partial_json(json).unwrap();
        assert!(result.is_some());
        let value = result.unwrap();
        assert_eq!(value["text"], "use { for scope and } to close");
        assert_eq!(value["count"], 5);
    }

    #[test]
    fn test_escaped_quotes_with_braces() {
        let partial = r#"{"code": "fn main() { println!(\"hello\"); }"#;
        let result = try_parse_partial_json(partial).unwrap();
        assert!(result.is_some());
        let value = result.unwrap();
        assert_eq!(value["code"], "fn main() { println!(\"hello\"); }");
    }

    #[test]
    fn test_deeply_nested_partial_json_3_levels() {
        let partial = r#"{"level1": {"level2": {"level3": {"name": "deep""#;
        let result = try_parse_partial_json(partial).unwrap();
        assert!(result.is_some());
        let value = result.unwrap();
        assert_eq!(value["level1"]["level2"]["level3"]["name"], "deep");
    }

    #[test]
    fn test_deeply_nested_partial_json_4_levels() {
        let partial = r#"{"a": {"b": {"c": {"d": {"value": 42"#;
        let result = try_parse_partial_json(partial).unwrap();
        assert!(result.is_some());
        let value = result.unwrap();
        assert_eq!(value["a"]["b"]["c"]["d"]["value"], 42);
    }

    #[test]
    fn test_deeply_nested_with_arrays() {
        let partial = r#"{"data": {"items": [{"name": "first"}, {"name": "second"#;
        let result = try_parse_partial_json(partial).unwrap();
        assert!(result.is_some());
        let value = result.unwrap();
        assert_eq!(value["data"]["items"][0]["name"], "first");
    }

    #[test]
    fn test_deeply_nested_complete() {
        let json = r#"{"a": {"b": {"c": {"d": {"e": "five"}}}}}"#;
        let result = try_parse_partial_json(json).unwrap();
        assert!(result.is_some());
        let value = result.unwrap();
        assert_eq!(value["a"]["b"]["c"]["d"]["e"], "five");
    }

    #[test]
    fn test_unicode_escape_sequences_complete() {
        let json = r#"{"name": "\u0048\u0065\u006c\u006c\u006f"}"#;
        let result = try_parse_partial_json(json).unwrap();
        assert!(result.is_some());
        let value = result.unwrap();
        assert_eq!(value["name"], "Hello");
    }

    #[test]
    fn test_unicode_escape_sequences_partial() {
        let partial = r#"{"name": "\u0048\u0065\u006c\u006c\u006f""#;
        let result = try_parse_partial_json(partial).unwrap();
        assert!(result.is_some());
        let value = result.unwrap();
        assert_eq!(value["name"], "Hello");
    }

    #[test]
    fn test_unicode_escape_incomplete_sequence() {
        // Incomplete unicode escape sequences cannot be parsed - return None
        let partial = r#"{"name": "\u004"#;
        let result = try_parse_partial_json(partial).unwrap();
        // Cannot parse incomplete unicode escape - need more data
        assert!(result.is_none());
    }

    #[test]
    fn test_unicode_escape_with_other_fields() {
        let partial = r#"{"greeting": "\u0048\u0069", "count": 42"#;
        let result = try_parse_partial_json(partial).unwrap();
        assert!(result.is_some());
        let value = result.unwrap();
        assert_eq!(value["greeting"], "Hi");
        assert_eq!(value["count"], 42);
    }

    #[test]
    fn test_mixed_unicode_and_regular_text() {
        let json = r#"{"text": "Say \u0048\u0069 to everyone"}"#;
        let result = try_parse_partial_json(json).unwrap();
        assert!(result.is_some());
        let value = result.unwrap();
        assert_eq!(value["text"], "Say Hi to everyone");
    }

    #[test]
    fn test_literal_newline_in_string() {
        let json = "{\"message\": \"Hello\nWorld\"}";
        let result = try_parse_partial_json(json).unwrap();
        assert!(result.is_some());
        let value = result.unwrap();
        assert_eq!(value["message"], "Hello\nWorld");
    }

    #[test]
    fn test_literal_tab_in_string() {
        let json = "{\"message\": \"Hello\tWorld\"}";
        let result = try_parse_partial_json(json).unwrap();
        assert!(result.is_some());
        let value = result.unwrap();
        assert_eq!(value["message"], "Hello\tWorld");
    }

    #[test]
    fn test_multiple_literal_newlines() {
        let json = "{\"message\": \"Line 1\nLine 2\nLine 3\"}";
        let result = try_parse_partial_json(json).unwrap();
        assert!(result.is_some());
        let value = result.unwrap();
        assert_eq!(value["message"], "Line 1\nLine 2\nLine 3");
    }

    #[test]
    fn test_literal_crlf_in_string() {
        let json = "{\"message\": \"Hello\r\nWorld\"}";
        let result = try_parse_partial_json(json).unwrap();
        assert!(result.is_some());
        let value = result.unwrap();
        assert_eq!(value["message"], "Hello\r\nWorld");
    }

    #[test]
    fn test_escape_control_chars_preserves_escaped() {
        let input = r#"{"text": "already escaped\\n here"}"#;
        let escaped = escape_control_chars_in_strings(input);
        assert_eq!(escaped, input);
    }

    #[test]
    fn test_escape_control_chars_handles_mixed() {
        let input = "{\"text\": \"literal\nnewline and escaped\\n too\"}";
        let escaped = escape_control_chars_in_strings(input);
        assert_eq!(escaped, "{\"text\": \"literal\\nnewline and escaped\\n too\"}");
    }

    // ============================================================
    // Regression tests for curly/smart quotes in streaming parsing
    // These ensure the streaming path handles LLM outputs correctly
    // ============================================================

    #[test]
    fn test_curly_quotes_in_json_structure() {
        // LLM uses curly quotes for JSON structure (outside strings)
        // Input: { "tool": "Bash" }
        let input = "{ \u{201C}tool\u{201D}: \u{201C}Bash\u{201D} }";
        let result = try_parse_partial_json(input).unwrap();
        assert!(result.is_some(), "Should parse JSON with curly quotes as structure");
        let value = result.unwrap();
        assert_eq!(value["tool"], "Bash");
    }

    #[test]
    fn test_curly_quotes_inside_string_value() {
        // LLM uses curly quotes inside a string value (should be escaped)
        // Input: {"message": "blend of "cat" and "ethos""}
        let input = "{\"message\": \"blend of \u{201C}cat\u{201D} and \u{201C}ethos\u{201D}\"}";
        let result = try_parse_partial_json(input).unwrap();
        assert!(result.is_some(), "Should parse JSON with curly quotes inside strings");
        let value = result.unwrap();
        let msg = value["message"].as_str().unwrap();
        assert!(msg.contains("cat"), "Message should contain 'cat'");
        assert!(msg.contains("ethos"), "Message should contain 'ethos'");
    }

    #[test]
    fn test_escaped_quotes_in_string_value() {
        // LLM correctly escapes quotes in string (common case)
        let input = r#"{"message": "Your username is \"catethos\". It means \"cat\" + \"ethos\"."}"#;
        let result = try_parse_partial_json(input).unwrap();
        assert!(result.is_some(), "Should parse JSON with escaped quotes");
        let value = result.unwrap();
        let msg = value["message"].as_str().unwrap();
        assert!(msg.contains("catethos"), "Message should contain 'catethos'");
        assert!(msg.contains("cat"), "Message should contain 'cat'");
    }

    #[test]
    fn test_final_response_with_escaped_quotes_and_special_chars() {
        // Simulates the exact FinalResponse format that was failing
        let input = r#"{
  "tool": "FinalResponse",
  "message": "Your current username is \"catethos\". It appears to be a playful blend of two words: \"cat\" and \"ethos.\" \"Cat\" evokes the image of a curious, independent feline, while \"ethos\" refers to the characteristic spirit, values, or beliefs of a community or individual."
}"#;
        let result = try_parse_partial_json(input).unwrap();
        assert!(result.is_some(), "Should parse FinalResponse with escaped quotes");
        let value = result.unwrap();
        assert_eq!(value["tool"], "FinalResponse");
        let msg = value["message"].as_str().unwrap();
        assert!(msg.contains("catethos"));
        assert!(msg.contains("cat"));
        assert!(msg.contains("ethos"));
    }

    #[test]
    fn test_mixed_curly_and_escaped_quotes() {
        // LLM mixes curly quotes (for emphasis) with proper escaped quotes
        // Input: {"msg": "He said \"hello\" and used "emphasis" marks"}
        let input = "{\"msg\": \"He said \\\"hello\\\" and used \u{201C}emphasis\u{201D} marks\"}";
        let result = try_parse_partial_json(input).unwrap();
        assert!(result.is_some(), "Should parse mixed quote styles");
        let value = result.unwrap();
        let msg = value["msg"].as_str().unwrap();
        assert!(msg.contains("hello"));
        assert!(msg.contains("emphasis"));
    }

    #[test]
    fn test_curly_quotes_with_newlines_in_string() {
        // Combines both issues: curly quotes AND literal newlines
        let input = "{\"message\": \"First line\nSecond with \u{201C}quotes\u{201D}\"}";
        let result = try_parse_partial_json(input).unwrap();
        assert!(result.is_some(), "Should parse curly quotes with newlines");
        let value = result.unwrap();
        let msg = value["message"].as_str().unwrap();
        assert!(msg.contains("First line"));
        assert!(msg.contains("quotes"));
    }

    #[test]
    fn test_normalize_quotes_preserves_already_escaped() {
        // Already escaped quotes should not be double-escaped
        let input = r#"{"text": "already \"escaped\" here"}"#;
        let normalized = normalize_quotes(input);
        assert_eq!(normalized, input, "Already escaped quotes should be preserved");
    }

    #[test]
    fn test_normalize_quotes_handles_em_dash() {
        // Em-dash (—) is common in LLM outputs and should pass through
        let input = r#"{"message": "curiosity—much like a cat"}"#;
        let result = try_parse_partial_json(input).unwrap();
        assert!(result.is_some());
        let value = result.unwrap();
        assert!(value["message"].as_str().unwrap().contains("—"));
    }

    #[test]
    fn test_json_with_embedded_markdown_code_blocks() {
        // This is the exact failure case: LLM returns valid JSON, but the message field
        // contains embedded markdown code blocks (```json). The parser should NOT try
        // to extract from these embedded code blocks.
        let input = r#"{"tool":"FinalResponse","message":"Here's a JSON example:\n\n```json\n{\"data\": [1, 2, 3]}\n```\n\nThis shows how to format data."}"#;
        let result = try_parse_partial_json(input).unwrap();
        assert!(result.is_some(), "Should parse valid JSON containing embedded markdown");
        let value = result.unwrap();
        assert_eq!(value["tool"], "FinalResponse");
        let msg = value["message"].as_str().unwrap();
        assert!(msg.contains("```json"));
        assert!(msg.contains("```\n\nThis shows"));
    }

    #[test]
    fn test_json_with_multiple_embedded_code_blocks() {
        // Multiple code blocks in the message field
        let input = r#"{"tool":"FinalResponse","message":"Example 1:\n```python\nprint('hello')\n```\n\nExample 2:\n```json\n{\"x\": 1}\n```\nDone."}"#;
        let result = try_parse_partial_json(input).unwrap();
        assert!(result.is_some(), "Should parse JSON with multiple embedded code blocks");
        let value = result.unwrap();
        assert_eq!(value["tool"], "FinalResponse");
        let msg = value["message"].as_str().unwrap();
        assert!(msg.contains("```python"));
        assert!(msg.contains("```json"));
    }

    #[test]
    fn test_actual_markdown_wrapped_json_still_works() {
        // When the LLM actually wraps JSON in markdown, extraction should still work
        let input = "```json\n{\"tool\": \"FinalResponse\", \"message\": \"Hello\"}\n```";
        let result = try_parse_partial_json(input).unwrap();
        assert!(result.is_some(), "Should extract JSON from actual markdown wrapper");
        let value = result.unwrap();
        assert_eq!(value["tool"], "FinalResponse");
        assert_eq!(value["message"], "Hello");
    }
}