simplify_baml 0.2.0

/// JSON parser with type coercion
///
/// This module implements lenient JSON parsing and type coercion.
/// It attempts to parse LLM outputs that may not be perfectly formatted JSON,
/// and coerces the values to match the expected BAML types.

use anyhow::{Context, Result};
use serde_json::Value as JsonValue;
use std::collections::HashMap;

use crate::ir::{BamlValue, FieldType, IR};

/// 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
}

/// Find the bounds of a JSON object using string-aware brace matching
fn find_json_object_bounds(text: &str) -> Option<(usize, usize)> {
    find_json_bounds(text, '{', '}')
}

/// Find the bounds of a JSON array using string-aware bracket matching
fn find_json_array_bounds(text: &str) -> Option<(usize, usize)> {
    find_json_bounds(text, '[', ']')
}

/// Find matching JSON structure bounds, properly handling strings
/// Returns byte indices suitable for string slicing
fn find_json_bounds(text: &str, open_char: char, close_char: char) -> Option<(usize, usize)> {
    let mut start_pos = None;
    let mut depth = 0;
    let mut in_string = false;
    let mut escape_next = false;

    for (byte_idx, c) in text.char_indices() {
        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;
        }

        if c == open_char {
            if start_pos.is_none() {
                start_pos = Some(byte_idx);
            }
            depth += 1;
        } else if c == close_char {
            depth -= 1;
            if depth == 0 && start_pos.is_some() {
                // Return byte index of the closing char + its byte length
                return Some((start_pos.unwrap(), byte_idx + c.len_utf8() - 1));
            }
        }
    }

    None
}

pub struct Parser<'a> {
    ir: &'a IR,
}

impl<'a> Parser<'a> {
    pub fn new(ir: &'a IR) -> Self {
        Self { ir }
    }

    /// Parse and coerce a raw LLM response string to a BamlValue
    ///
    /// # Arguments
    /// * `raw_response` - The raw string from the LLM
    /// * `target_type` - The expected output type
    ///
    /// # Returns
    /// The parsed and coerced BamlValue
    pub fn parse(&self, raw_response: &str, target_type: &FieldType) -> Result<BamlValue> {
        // Step 0: Normalize curly/smart quotes to ASCII quotes
        // LLMs sometimes output Unicode quotes which break JSON parsing
        let normalized = normalize_quotes(raw_response);

        // Step 1: Extract JSON from the response (lenient extraction)
        let json_str = self.extract_json(&normalized)?;

        // Step 2: Parse JSON
        let json_value: JsonValue = serde_json::from_str(&json_str)
            .context("Failed to parse JSON from LLM response")?;

        // Step 3: Coerce to target type
        self.coerce(&json_value, target_type)
    }

    /// Extract JSON from a response that may contain markdown code blocks or extra text
    ///
    /// NOTE: We first check if the response is already valid JSON before attempting
    /// markdown extraction. This prevents issues where valid JSON containing embedded
    /// markdown code blocks (e.g., in a string field) would be incorrectly extracted.
    fn extract_json(&self, response: &str) -> Result<String> {
        let response = response.trim();

        // First, try to parse the response as-is to see if it's already valid JSON
        // This handles the common case where LLM returns raw JSON without markdown
        if serde_json::from_str::<serde_json::Value>(response).is_ok() {
            return Ok(response.to_string());
        }

        // Try to find JSON boundaries using string-aware matching
        // This handles cases where there's extra text before/after the JSON
        let array_bounds = find_json_array_bounds(response);
        let object_bounds = find_json_object_bounds(response);

        // Return whichever starts first (to handle arrays containing objects correctly)
        match (array_bounds, object_bounds) {
            (Some((arr_start, arr_end)), Some((obj_start, obj_end))) => {
                let extracted = if arr_start <= obj_start {
                    &response[arr_start..=arr_end]
                } else {
                    &response[obj_start..=obj_end]
                };
                // Verify extracted content is valid JSON before returning
                if serde_json::from_str::<serde_json::Value>(extracted).is_ok() {
                    return Ok(extracted.to_string());
                }
            }
            (Some((start, end)), None) => {
                let extracted = &response[start..=end];
                if serde_json::from_str::<serde_json::Value>(extracted).is_ok() {
                    return Ok(extracted.to_string());
                }
            }
            (None, Some((start, end))) => {
                let extracted = &response[start..=end];
                if serde_json::from_str::<serde_json::Value>(extracted).is_ok() {
                    return Ok(extracted.to_string());
                }
            }
            (None, None) => {}
        }

        // Only now try markdown extraction as a fallback
        // This handles responses wrapped in ```json code blocks
        let response_lower = response.to_lowercase();
        if let Some(start) = response_lower.find("```json") {
            let json_start = start + 7; // len("```json")
            // Find the closing ``` after the opening one
            if let Some(end_offset) = response[json_start..].find("```") {
                let json_end = json_start + end_offset;
                return Ok(response[json_start..json_end].trim().to_string());
            }
        }

        // Check for code block without language specifier
        if let Some(start) = response.find("```") {
            if let Some(end) = response[start + 3..].find("```") {
                let json_start = start + 3;
                let json_end = start + 3 + end;
                let content = response[json_start..json_end].trim();
                // Only use this if it looks like JSON
                if content.starts_with('{') || content.starts_with('[') {
                    return Ok(content.to_string());
                }
            }
        }

        // If nothing found, assume the whole response is JSON
        Ok(response.to_string())
    }

    /// Coerce a JSON value to match the target BAML type
    fn coerce(&self, value: &JsonValue, target_type: &FieldType) -> Result<BamlValue> {
        match target_type {
            FieldType::String => self.coerce_string(value),
            FieldType::Int => self.coerce_int(value),
            FieldType::Float => self.coerce_float(value),
            FieldType::Bool => self.coerce_bool(value),
            FieldType::Enum(enum_name) => self.coerce_enum(value, enum_name),
            FieldType::Class(class_name) => {
                // Try as class first, then fall back to enum
                // This handles the case where macros generate Class for all custom types
                if self.ir.find_class(class_name).is_some() {
                    self.coerce_class(value, class_name)
                } else if self.ir.find_enum(class_name).is_some() {
                    self.coerce_enum(value, class_name)
                } else {
                    anyhow::bail!("Type '{}' not found (neither class nor enum)", class_name)
                }
            }
            FieldType::List(inner) => self.coerce_list(value, inner),
            FieldType::Map(k, v) => self.coerce_map(value, k, v),
            FieldType::Union(types) => self.coerce_union(value, types),
            FieldType::TaggedEnum(name) => self.coerce_tagged_enum(value, name),
        }
    }

    fn coerce_string(&self, value: &JsonValue) -> Result<BamlValue> {
        match value {
            JsonValue::String(s) => Ok(BamlValue::String(s.clone())),
            JsonValue::Number(n) => Ok(BamlValue::String(n.to_string())),
            JsonValue::Bool(b) => Ok(BamlValue::String(b.to_string())),
            JsonValue::Null => Ok(BamlValue::String("".to_string())),
            JsonValue::Object(obj) => {
                // LLM may wrap the value in an object like { "value": "text" } or { "string": "text" }
                // Try common field names
                for field_name in ["value", "Value", "string", "String", "text", "Text", "result", "Result"] {
                    if let Some(inner) = obj.get(field_name) {
                        return self.coerce_string(inner);
                    }
                }
                // If object has only one field, use that
                if obj.len() == 1 {
                    if let Some((_, inner)) = obj.iter().next() {
                        return self.coerce_string(inner);
                    }
                }
                anyhow::bail!("Cannot coerce object to string: {:?}", value)
            }
            _ => anyhow::bail!("Cannot coerce {:?} to string", value),
        }
    }

    fn coerce_int(&self, value: &JsonValue) -> Result<BamlValue> {
        match value {
            JsonValue::Number(n) => {
                if let Some(i) = n.as_i64() {
                    Ok(BamlValue::Int(i))
                } else if let Some(f) = n.as_f64() {
                    if f.fract() != 0.0 {
                        anyhow::bail!("Cannot coerce non-integral float {} to int", f)
                    }
                    if f > i64::MAX as f64 || f < i64::MIN as f64 {
                        anyhow::bail!("Float {} overflows i64 range", f)
                    }
                    Ok(BamlValue::Int(f as i64))
                } else {
                    anyhow::bail!("Cannot coerce number to int")
                }
            }
            JsonValue::String(s) => {
                let i = s.parse::<i64>()
                    .context("Cannot parse string as int")?;
                Ok(BamlValue::Int(i))
            }
            JsonValue::Object(obj) => {
                // LLM may wrap the value in an object like { "value": 42 } or { "int": 42 }
                // Try common field names
                for field_name in ["value", "Value", "int", "Int", "number", "Number", "result", "Result"] {
                    if let Some(inner) = obj.get(field_name) {
                        return self.coerce_int(inner);
                    }
                }
                // If object has only one field, use that
                if obj.len() == 1 {
                    if let Some((_, inner)) = obj.iter().next() {
                        return self.coerce_int(inner);
                    }
                }
                anyhow::bail!("Cannot coerce object to int: {:?}", value)
            }
            _ => anyhow::bail!("Cannot coerce {:?} to int", value),
        }
    }

    fn coerce_float(&self, value: &JsonValue) -> Result<BamlValue> {
        match value {
            JsonValue::Number(n) => {
                if let Some(f) = n.as_f64() {
                    Ok(BamlValue::Float(f))
                } else {
                    anyhow::bail!("Cannot coerce number to float")
                }
            }
            JsonValue::String(s) => {
                let f = s.parse::<f64>()
                    .context("Cannot parse string as float")?;
                Ok(BamlValue::Float(f))
            }
            JsonValue::Object(obj) => {
                // LLM may wrap the value in an object like { "value": 3.14 } or { "float": 3.14 }
                // Try common field names
                for field_name in ["value", "Value", "float", "Float", "number", "Number", "result", "Result"] {
                    if let Some(inner) = obj.get(field_name) {
                        return self.coerce_float(inner);
                    }
                }
                // If object has only one field, use that
                if obj.len() == 1 {
                    if let Some((_, inner)) = obj.iter().next() {
                        return self.coerce_float(inner);
                    }
                }
                anyhow::bail!("Cannot coerce object to float: {:?}", value)
            }
            _ => anyhow::bail!("Cannot coerce {:?} to float", value),
        }
    }

    fn coerce_bool(&self, value: &JsonValue) -> Result<BamlValue> {
        match value {
            JsonValue::Bool(b) => Ok(BamlValue::Bool(*b)),
            JsonValue::String(s) => {
                let s_lower = s.to_lowercase();
                if s_lower == "true" || s_lower == "yes" || s_lower == "1" {
                    Ok(BamlValue::Bool(true))
                } else if s_lower == "false" || s_lower == "no" || s_lower == "0" {
                    Ok(BamlValue::Bool(false))
                } else {
                    anyhow::bail!("Cannot parse '{}' as bool", s)
                }
            }
            JsonValue::Number(n) => {
                if let Some(i) = n.as_i64() {
                    Ok(BamlValue::Bool(i != 0))
                } else {
                    anyhow::bail!("Cannot coerce number to bool")
                }
            }
            JsonValue::Object(obj) => {
                // LLM may wrap the value in an object like { "value": true } or { "bool": true }
                // Try common field names
                for field_name in ["value", "Value", "bool", "Bool", "result", "Result"] {
                    if let Some(inner) = obj.get(field_name) {
                        return self.coerce_bool(inner);
                    }
                }
                // If object has only one field, use that
                if obj.len() == 1 {
                    if let Some((_, inner)) = obj.iter().next() {
                        return self.coerce_bool(inner);
                    }
                }
                anyhow::bail!("Cannot coerce object to bool: {:?}", value)
            }
            _ => anyhow::bail!("Cannot coerce {:?} to bool", value),
        }
    }

    fn coerce_enum(&self, value: &JsonValue, enum_name: &str) -> Result<BamlValue> {
        let e = self.ir.find_enum(enum_name)
            .ok_or_else(|| anyhow::anyhow!("Enum '{}' not found", enum_name))?;

        let str_value = match value {
            JsonValue::String(s) => s.clone(),
            JsonValue::Number(n) => n.to_string(),
            JsonValue::Bool(b) => b.to_string(),
            JsonValue::Null => anyhow::bail!("Cannot coerce null to enum '{}'", enum_name),
            JsonValue::Array(_) => anyhow::bail!("Cannot coerce array to enum '{}': arrays are not valid enum values", enum_name),
            JsonValue::Object(_) => anyhow::bail!("Cannot coerce object to enum '{}': objects are not valid enum values", enum_name),
        };

        // Check if the value is a valid enum variant
        if e.values.contains(&str_value) {
            Ok(BamlValue::String(str_value))
        } else {
            // Try case-insensitive match
            let lower = str_value.to_lowercase();
            for variant in &e.values {
                if variant.to_lowercase() == lower {
                    return Ok(BamlValue::String(variant.clone()));
                }
            }
            anyhow::bail!("'{}' is not a valid variant of enum '{}'", str_value, enum_name)
        }
    }

    fn coerce_tagged_enum(&self, value: &JsonValue, enum_name: &str) -> Result<BamlValue> {
        let te = self.ir.find_tagged_enum(enum_name)
            .ok_or_else(|| anyhow::anyhow!("Tagged enum '{}' not found", enum_name))?;

        let obj = value.as_object()
            .ok_or_else(|| anyhow::anyhow!("Expected object for tagged enum '{}'", enum_name))?;

        // Get the tag value
        let tag_value = obj.get(&te.tag_field)
            .ok_or_else(|| anyhow::anyhow!("Missing tag field '{}' for tagged enum '{}'", te.tag_field, enum_name))?;

        let tag_str = tag_value.as_str()
            .ok_or_else(|| anyhow::anyhow!("Tag field '{}' must be a string", te.tag_field))?;

        // Find matching variant (case-insensitive)
        let variant = te.variants.iter()
            .find(|v| v.name.eq_ignore_ascii_case(tag_str))
            .ok_or_else(|| anyhow::anyhow!("'{}' is not a valid variant of tagged enum '{}'", tag_str, enum_name))?;

        // Build result map with the normalized tag and all variant fields
        let mut result = HashMap::new();
        result.insert(te.tag_field.clone(), BamlValue::String(variant.name.clone()));

        for field in &variant.fields {
            if let Some(field_value) = obj.get(&field.name) {
                if field_value.is_null() {
                    if !field.optional {
                        anyhow::bail!(
                            "Field '{}' in variant '{}' is required but got null",
                            field.name,
                            variant.name
                        );
                    }
                } else {
                    let coerced = self.coerce(field_value, &field.field_type)?;
                    result.insert(field.name.clone(), coerced);
                }
            } else if !field.optional {
                anyhow::bail!("Missing required field '{}' in variant '{}'", field.name, variant.name);
            }
        }

        Ok(BamlValue::Map(result))
    }

    fn coerce_class(&self, value: &JsonValue, class_name: &str) -> Result<BamlValue> {
        let class = self.ir.find_class(class_name)
            .ok_or_else(|| anyhow::anyhow!("Class '{}' not found", class_name))?;

        let obj = value.as_object()
            .ok_or_else(|| anyhow::anyhow!("Expected object for class '{}'", class_name))?;

        let mut result = HashMap::new();

        for field in &class.fields {
            if let Some(field_value) = obj.get(&field.name) {
                if field_value.is_null() {
                    if !field.optional {
                        anyhow::bail!(
                            "Field '{}' in class '{}' is required but got null",
                            field.name,
                            class_name
                        );
                    }
                    // Skip null values for optional fields (don't add to result)
                } else {
                    let coerced = self.coerce(field_value, &field.field_type)?;
                    result.insert(field.name.clone(), coerced);
                }
            } else if !field.optional {
                anyhow::bail!("Missing required field '{}' in class '{}'", field.name, class_name);
            }
        }

        Ok(BamlValue::Map(result))
    }

    fn coerce_list(&self, value: &JsonValue, inner_type: &FieldType) -> Result<BamlValue> {
        if let Some(arr) = value.as_array() {
            let coerced: Result<Vec<BamlValue>> = arr.iter()
                .map(|item| self.coerce(item, inner_type))
                .collect();
            Ok(BamlValue::List(coerced?))
        } else {
            // Scalar-to-list coercion: wrap non-array value in singleton list
            let coerced = self.coerce(value, inner_type)
                .context("Failed to coerce scalar to list element")?;
            Ok(BamlValue::List(vec![coerced]))
        }
    }

    fn coerce_map(&self, value: &JsonValue, key_type: &FieldType, value_type: &FieldType) -> Result<BamlValue> {
        let obj = value.as_object()
            .ok_or_else(|| anyhow::anyhow!("Expected object for map"))?;

        self.validate_map_key_type(key_type)?;

        let coerced: Result<HashMap<String, BamlValue>> = obj.iter()
            .map(|(k, v)| {
                self.coerce_map_key(k, key_type)?;
                self.coerce(v, value_type)
                    .map(|coerced_v| (k.clone(), coerced_v))
            })
            .collect();

        Ok(BamlValue::Map(coerced?))
    }

    fn validate_map_key_type(&self, key_type: &FieldType) -> Result<()> {
        match key_type {
            FieldType::String | FieldType::Int => Ok(()),
            _ => anyhow::bail!(
                "Unsupported map key type: {:?}. Only String and Int keys are supported.",
                key_type
            ),
        }
    }

    fn coerce_map_key(&self, key: &str, key_type: &FieldType) -> Result<()> {
        match key_type {
            FieldType::String => Ok(()),
            FieldType::Int => {
                key.parse::<i64>()
                    .map_err(|_| anyhow::anyhow!("Map key '{}' cannot be parsed as integer", key))?;
                Ok(())
            }
            _ => anyhow::bail!("Unsupported map key type: {:?}", key_type),
        }
    }

    fn coerce_union(&self, value: &JsonValue, types: &[FieldType]) -> Result<BamlValue> {
        // Try each type in order until one succeeds
        for t in types {
            if let Ok(coerced) = self.coerce(value, t) {
                return Ok(coerced);
            }
        }
        anyhow::bail!("Cannot coerce {:?} to any of the union types", value)
    }
}

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

    #[test]
    fn test_extract_json_from_markdown() {
        let ir = IR::new();
        let parser = Parser::new(&ir);

        let response = r#"
        Here's the result:
        ```json
        {"name": "John", "age": 30}
        ```
        "#;

        let json = parser.extract_json(response).unwrap();
        assert_eq!(json.trim(), r#"{"name": "John", "age": 30}"#);
    }

    #[test]
    fn test_extract_json_from_uppercase_markdown() {
        let ir = IR::new();
        let parser = Parser::new(&ir);

        let response = r#"
        Here's the result:
        ```JSON
        {"name": "John", "age": 30}
        ```
        "#;

        let json = parser.extract_json(response).unwrap();
        assert_eq!(json.trim(), r#"{"name": "John", "age": 30}"#);
    }

    #[test]
    fn test_extract_json_from_mixed_case_markdown() {
        let ir = IR::new();
        let parser = Parser::new(&ir);

        let response = r#"
        ```Json
        {"name": "Alice"}
        ```
        "#;

        let json = parser.extract_json(response).unwrap();
        assert_eq!(json.trim(), r#"{"name": "Alice"}"#);
    }

    #[test]
    fn test_coerce_int_from_string() {
        let ir = IR::new();
        let parser = Parser::new(&ir);

        let value = JsonValue::String("42".to_string());
        let result = parser.coerce_int(&value).unwrap();

        assert_eq!(result.as_int(), Some(42));
    }

    #[test]
    fn test_parse_class() {
        let mut ir = IR::new();
        ir.classes.push(Class {
            name: "Person".to_string(),
            description: None,
            fields: vec![
                Field {
                    name: "name".to_string(),
                    field_type: FieldType::String,
                    optional: false,
                description: None,
                },
                Field {
                    name: "age".to_string(),
                    field_type: FieldType::Int,
                    optional: false,
                description: None,
                },
            ],
        });

        let parser = Parser::new(&ir);
        let response = r#"{"name": "John", "age": 30}"#;

        let result = parser.parse(response, &FieldType::Class("Person".to_string())).unwrap();

        if let BamlValue::Map(map) = result {
            assert_eq!(map.get("name").and_then(|v| v.as_string()), Some("John"));
            assert_eq!(map.get("age").and_then(|v| v.as_int()), Some(30));
        } else {
            panic!("Expected Map");
        }
    }

    #[test]
    fn test_extract_json_with_braces_in_string() {
        let ir = IR::new();
        let parser = Parser::new(&ir);

        let response = r#"{"text": "use { for scope and } to close"}"#;
        let json = parser.extract_json(response).unwrap();
        assert_eq!(json, r#"{"text": "use { for scope and } to close"}"#);

        let parsed: JsonValue = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed["text"], "use { for scope and } to close");
    }

    #[test]
    fn test_extract_json_with_brackets_in_string() {
        let ir = IR::new();
        let parser = Parser::new(&ir);

        let response = r#"{"code": "let arr = [1, 2, 3];"}"#;
        let json = parser.extract_json(response).unwrap();
        assert_eq!(json, r#"{"code": "let arr = [1, 2, 3];"}"#);

        let parsed: JsonValue = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed["code"], "let arr = [1, 2, 3];");
    }

    #[test]
    fn test_extract_json_with_nested_braces_in_string() {
        let ir = IR::new();
        let parser = Parser::new(&ir);

        let response = r#"Here's the result: {"message": "JSON example: {\"key\": \"value\"}"}"#;
        let json = parser.extract_json(response).unwrap();
        
        let parsed: JsonValue = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed["message"], r#"JSON example: {"key": "value"}"#);
    }

    #[test]
    fn test_extract_json_complex_string_content() {
        let ir = IR::new();
        let parser = Parser::new(&ir);

        let response = r#"Output: {"code": "fn main() { println!(\"hello\"); }", "lang": "rust"}"#;
        let json = parser.extract_json(response).unwrap();
        
        let parsed: JsonValue = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed["code"], "fn main() { println!(\"hello\"); }");
        assert_eq!(parsed["lang"], "rust");
    }

    #[test]
    fn test_extract_json_array() {
        let ir = IR::new();
        let parser = Parser::new(&ir);

        let response = r#"[{"name": "Regina"}, {"name": "Guaguaxuan"}]"#;
        let json = parser.extract_json(response).unwrap();
        println!("Extracted JSON: {}", json);
        
        let parsed: JsonValue = serde_json::from_str(&json).unwrap();
        assert!(parsed.is_array());
        assert_eq!(parsed.as_array().unwrap().len(), 2);
    }

    #[test]
    fn test_coerce_int_from_integral_float() {
        let ir = IR::new();
        let parser = Parser::new(&ir);

        let value: JsonValue = serde_json::from_str("3.0").unwrap();
        let result = parser.coerce_int(&value).unwrap();
        assert_eq!(result.as_int(), Some(3));
    }

    #[test]
    fn test_coerce_int_rejects_non_integral_float() {
        let ir = IR::new();
        let parser = Parser::new(&ir);

        let value: JsonValue = serde_json::from_str("3.5").unwrap();
        let result = parser.coerce_int(&value);
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("non-integral"));
    }

    #[test]
    fn test_coerce_int_rejects_overflow() {
        let ir = IR::new();
        let parser = Parser::new(&ir);

        let large_float = (i64::MAX as f64) * 2.0;
        let value: JsonValue = serde_json::from_str(&format!("{:.0}", large_float)).unwrap();
        let result = parser.coerce_int(&value);
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("overflow"));
    }

    #[test]
    fn test_coerce_enum_rejects_array() {
        let mut ir = IR::new();
        ir.enums.push(Enum {
            name: "Status".to_string(),
            values: vec!["Active".to_string(), "Inactive".to_string()],
            description: None,
        });

        let parser = Parser::new(&ir);
        let value: JsonValue = serde_json::from_str(r#"[1, 2]"#).unwrap();
        let result = parser.coerce_enum(&value, "Status");
        
        assert!(result.is_err());
        let err_msg = result.unwrap_err().to_string();
        assert!(err_msg.contains("array"), "Error should mention array: {}", err_msg);
    }

    #[test]
    fn test_coerce_enum_rejects_object() {
        let mut ir = IR::new();
        ir.enums.push(Enum {
            name: "Status".to_string(),
            values: vec!["Active".to_string(), "Inactive".to_string()],
            description: None,
        });

        let parser = Parser::new(&ir);
        let value: JsonValue = serde_json::from_str(r#"{"key": "value"}"#).unwrap();
        let result = parser.coerce_enum(&value, "Status");
        
        assert!(result.is_err());
        let err_msg = result.unwrap_err().to_string();
        assert!(err_msg.contains("object"), "Error should mention object: {}", err_msg);
    }

    #[test]
    fn test_coerce_enum_rejects_null() {
        let mut ir = IR::new();
        ir.enums.push(Enum {
            name: "Status".to_string(),
            values: vec!["Active".to_string(), "Inactive".to_string()],
            description: None,
        });

        let parser = Parser::new(&ir);
        let value = JsonValue::Null;
        let result = parser.coerce_enum(&value, "Status");
        
        assert!(result.is_err());
        let err_msg = result.unwrap_err().to_string();
        assert!(err_msg.contains("null"), "Error should mention null: {}", err_msg);
    }

    #[test]
    fn test_coerce_enum_accepts_valid_string() {
        let mut ir = IR::new();
        ir.enums.push(Enum {
            name: "Status".to_string(),
            values: vec!["Active".to_string(), "Inactive".to_string()],
            description: None,
        });

        let parser = Parser::new(&ir);
        let value = JsonValue::String("Active".to_string());
        let result = parser.coerce_enum(&value, "Status").unwrap();
        
        assert_eq!(result.as_string(), Some("Active"));
    }

    #[test]
    fn test_coerce_map_with_string_keys() {
        let ir = IR::new();
        let parser = Parser::new(&ir);

        let value: JsonValue = serde_json::from_str(r#"{"a": 1, "b": 2}"#).unwrap();
        let key_type = FieldType::String;
        let value_type = FieldType::Int;
        let result = parser.coerce_map(&value, &key_type, &value_type).unwrap();
        
        if let BamlValue::Map(map) = result {
            assert_eq!(map.len(), 2);
            assert_eq!(map.get("a").unwrap().as_int(), Some(1));
            assert_eq!(map.get("b").unwrap().as_int(), Some(2));
        } else {
            panic!("Expected Map");
        }
    }

    #[test]
    fn test_coerce_map_with_int_keys_valid() {
        let ir = IR::new();
        let parser = Parser::new(&ir);

        let value: JsonValue = serde_json::from_str(r#"{"1": "one", "2": "two"}"#).unwrap();
        let key_type = FieldType::Int;
        let value_type = FieldType::String;
        let result = parser.coerce_map(&value, &key_type, &value_type).unwrap();
        
        if let BamlValue::Map(map) = result {
            assert_eq!(map.len(), 2);
            assert_eq!(map.get("1").unwrap().as_string(), Some("one"));
            assert_eq!(map.get("2").unwrap().as_string(), Some("two"));
        } else {
            panic!("Expected Map");
        }
    }

    #[test]
    fn test_coerce_map_with_int_keys_invalid() {
        let ir = IR::new();
        let parser = Parser::new(&ir);

        let value: JsonValue = serde_json::from_str(r#"{"not_a_number": 1}"#).unwrap();
        let key_type = FieldType::Int;
        let value_type = FieldType::Int;
        let result = parser.coerce_map(&value, &key_type, &value_type);
        
        assert!(result.is_err());
        let err_msg = result.unwrap_err().to_string();
        assert!(err_msg.contains("not_a_number"), "Error should mention the invalid key: {}", err_msg);
        assert!(err_msg.contains("integer"), "Error should mention integer: {}", err_msg);
    }

    #[test]
    fn test_coerce_map_rejects_unsupported_key_type() {
        let ir = IR::new();
        let parser = Parser::new(&ir);

        let value: JsonValue = serde_json::from_str(r#"{"a": 1}"#).unwrap();
        let key_type = FieldType::Bool;
        let value_type = FieldType::Int;
        let result = parser.coerce_map(&value, &key_type, &value_type);
        
        assert!(result.is_err());
        let err_msg = result.unwrap_err().to_string();
        assert!(err_msg.contains("Unsupported map key type"), "Error should mention unsupported key type: {}", err_msg);
    }

    #[test]
    fn test_scalar_to_list_coercion_int() {
        let ir = IR::new();
        let parser = Parser::new(&ir);

        let value: JsonValue = serde_json::from_str("5").unwrap();
        let result = parser.coerce_list(&value, &FieldType::Int).unwrap();

        if let BamlValue::List(list) = result {
            assert_eq!(list.len(), 1);
            assert_eq!(list[0].as_int(), Some(5));
        } else {
            panic!("Expected List");
        }
    }

    #[test]
    fn test_scalar_to_list_coercion_string() {
        let ir = IR::new();
        let parser = Parser::new(&ir);

        let value = JsonValue::String("hello".to_string());
        let result = parser.coerce_list(&value, &FieldType::String).unwrap();

        if let BamlValue::List(list) = result {
            assert_eq!(list.len(), 1);
            assert_eq!(list[0].as_string(), Some("hello"));
        } else {
            panic!("Expected List");
        }
    }

    #[test]
    fn test_scalar_to_list_coercion_object_to_class_list() {
        let mut ir = IR::new();
        ir.classes.push(Class {
            name: "Item".to_string(),
            description: None,
            fields: vec![Field {
                name: "id".to_string(),
                field_type: FieldType::Int,
                optional: false,
                description: None,
            }],
        });

        let parser = Parser::new(&ir);
        let value: JsonValue = serde_json::from_str(r#"{"id": 42}"#).unwrap();
        let result = parser.coerce_list(&value, &FieldType::Class("Item".to_string())).unwrap();

        if let BamlValue::List(list) = result {
            assert_eq!(list.len(), 1);
            if let BamlValue::Map(map) = &list[0] {
                assert_eq!(map.get("id").unwrap().as_int(), Some(42));
            } else {
                panic!("Expected Map inside List");
            }
        } else {
            panic!("Expected List");
        }
    }

    #[test]
    fn test_scalar_to_list_coercion_fails_on_type_mismatch() {
        let ir = IR::new();
        let parser = Parser::new(&ir);

        let value = JsonValue::String("not_a_number".to_string());
        let result = parser.coerce_list(&value, &FieldType::Int);

        assert!(result.is_err());
    }

    #[test]
    fn test_empty_object_for_class_with_required_fields() {
        let mut ir = IR::new();
        ir.classes.push(Class {
            name: "Person".to_string(),
            description: None,
            fields: vec![
                Field {
                    name: "name".to_string(),
                    field_type: FieldType::String,
                    optional: false,
                    description: None,
                },
            ],
        });

        let parser = Parser::new(&ir);
        let value: JsonValue = serde_json::from_str(r#"{}"#).unwrap();
        let result = parser.coerce(&value, &FieldType::Class("Person".to_string()));

        assert!(result.is_err());
        let err_msg = result.unwrap_err().to_string();
        assert!(err_msg.contains("Missing required field"), "Error should mention missing field: {}", err_msg);
        assert!(err_msg.contains("name"), "Error should mention the field name: {}", err_msg);
    }

    #[test]
    fn test_empty_object_for_class_with_only_optional_fields() {
        let mut ir = IR::new();
        ir.classes.push(Class {
            name: "OptionalPerson".to_string(),
            description: None,
            fields: vec![
                Field {
                    name: "nickname".to_string(),
                    field_type: FieldType::String,
                    optional: true,
                    description: None,
                },
            ],
        });

        let parser = Parser::new(&ir);
        let value: JsonValue = serde_json::from_str(r#"{}"#).unwrap();
        let result = parser.coerce(&value, &FieldType::Class("OptionalPerson".to_string()));

        assert!(result.is_ok(), "Empty object should be valid when all fields are optional");
        if let BamlValue::Map(map) = result.unwrap() {
            assert!(map.is_empty(), "Map should be empty");
        } else {
            panic!("Expected Map");
        }
    }

    #[test]
    fn test_extract_json_with_extra_closing_braces() {
        let ir = IR::new();
        let parser = Parser::new(&ir);

        let response = r#"{"a": 1}}"#;
        let json = parser.extract_json(response).unwrap();
        assert_eq!(json, r#"{"a": 1}"#);

        let parsed: JsonValue = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed["a"], 1);
    }

    #[test]
    fn test_extract_json_with_multiple_extra_closing_braces() {
        let ir = IR::new();
        let parser = Parser::new(&ir);

        let response = r#"{"nested": {"b": 2}}}}}"#;
        let json = parser.extract_json(response).unwrap();
        assert_eq!(json, r#"{"nested": {"b": 2}}"#);

        let parsed: JsonValue = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed["nested"]["b"], 2);
    }

    #[test]
    fn test_extract_json_array_with_extra_closing_brackets() {
        let ir = IR::new();
        let parser = Parser::new(&ir);

        let response = r#"[1, 2, 3]]]"#;
        let json = parser.extract_json(response).unwrap();
        assert_eq!(json, r#"[1, 2, 3]"#);

        let parsed: JsonValue = serde_json::from_str(&json).unwrap();
        assert!(parsed.is_array());
        assert_eq!(parsed.as_array().unwrap().len(), 3);
    }

    #[test]
    fn test_curly_quotes_inside_string_are_escaped() {
        // Curly quotes inside strings should be escaped, not converted
        // Input: {"message": "blend of "cat" and "ethos""}
        let input = "{\"message\": \"blend of \u{201C}cat\u{201D} and \u{201C}ethos\u{201D}\"}";
        let normalized = normalize_quotes(input);
        // The curly quotes inside the string should become escaped quotes
        assert!(normalized.contains(r#"blend of \"cat\" and \"ethos\""#), "normalized was: {}", normalized);
        // Should be valid JSON
        let parsed: JsonValue = serde_json::from_str(&normalized).unwrap();
        assert!(parsed["message"].as_str().unwrap().contains("cat"));
    }

    #[test]
    fn test_curly_quotes_for_json_structure_are_converted() {
        // Curly quotes used for JSON structure should become ASCII quotes
        let input = "{ \u{201C}tool\u{201D}: \u{201C}Bash\u{201D} }";
        let normalized = normalize_quotes(input);
        assert_eq!(normalized, r#"{ "tool": "Bash" }"#);
    }

    #[test]
    fn test_literal_newlines_in_string_are_escaped() {
        let input = "{\"message\": \"Hello\nWorld\"}";
        let normalized = normalize_quotes(input);
        assert_eq!(normalized, r#"{"message": "Hello\nWorld"}"#);
        let parsed: JsonValue = serde_json::from_str(&normalized).unwrap();
        assert_eq!(parsed["message"], "Hello\nWorld");
    }

    #[test]
    fn test_final_response_exact_failure_case() {
        // This is the exact JSON that was failing in tao
        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. Put together, \"catethos\" suggests a personal philosophy that values curiosity, agility, and perhaps a touch of mischievous independence–much like a cat navigating the world on its own terms."
}"#;

        let ir = IR::new();
        let parser = Parser::new(&ir);

        // First check extract_json works
        let extracted = parser.extract_json(input);
        assert!(extracted.is_ok(), "extract_json failed: {:?}", extracted.err());

        // Check serde can parse it
        let json_str = extracted.unwrap();
        let parsed: Result<JsonValue, _> = serde_json::from_str(&json_str);
        assert!(parsed.is_ok(), "serde_json failed on: {}\nError: {:?}", json_str, parsed.err());

        let value = parsed.unwrap();
        assert_eq!(value["tool"], "FinalResponse");
        assert!(value["message"].as_str().unwrap().contains("catethos"));
    }

    #[test]
    fn test_json_with_embedded_markdown_code_blocks() {
        // This is the failure case: valid JSON with embedded markdown code blocks in a string field
        // The parser should NOT try to extract from these embedded code blocks
        let input = r#"{"tool": "FinalResponse", "message": "Here's a random rechart configuration in JSON format:\n\n```json\n{\n  \"data\": [\n    { \"name\": \"Jan\", \"sales\": 4000 }\n  ]\n}\n```\n\nThis represents a typical rechart configuration."}"#;

        let ir = IR::new();
        let parser = Parser::new(&ir);

        // extract_json should return the entire input since it's valid JSON
        let extracted = parser.extract_json(input);
        assert!(extracted.is_ok(), "extract_json failed: {:?}", extracted.err());

        let json_str = extracted.unwrap();
        let parsed: Result<JsonValue, _> = serde_json::from_str(&json_str);
        assert!(parsed.is_ok(), "serde_json failed: {:?}", parsed.err());

        let value = parsed.unwrap();
        assert_eq!(value["tool"], "FinalResponse");
        let msg = value["message"].as_str().unwrap();
        assert!(msg.contains("```json"), "message should contain the embedded code block");
        assert!(msg.contains("rechart configuration"), "message should contain the description");
    }

    #[test]
    fn test_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 ir = IR::new();
        let parser = Parser::new(&ir);

        let extracted = parser.extract_json(input);
        assert!(extracted.is_ok(), "extract_json failed: {:?}", extracted.err());

        let json_str = extracted.unwrap();
        let parsed: Result<JsonValue, _> = serde_json::from_str(&json_str);
        assert!(parsed.is_ok(), "serde_json failed: {:?}", parsed.err());

        let value = parsed.unwrap();
        assert_eq!(value["tool"], "FinalResponse");
        assert_eq!(value["message"], "Hello");
    }
}