/// Extract the first balanced JSON entity (`{...}` or `[...]`) from text.
///
/// Uses bracket counting with proper string/escape handling.
/// Returns the slice containing the JSON entity, or `None` if not found.
pub fn extract_json_entity(text: &str) -> Option<&str> {
    extract_all_json_entities(text).into_iter().next()
}

/// Extract all balanced JSON entities (`{...}` or `[...]`) from text.
///
/// Returns slices in the order they appear. Each entity is independently
/// balanced. Useful when an LLM returns multiple JSON objects in one response.
pub fn extract_all_json_entities(text: &str) -> Vec<&str> {
    let bytes = text.as_bytes();
    let len = bytes.len();
    let mut results = Vec::new();
    let mut i = 0;

    while i < len {
        match bytes[i] {
            b'{' | b'[' => {
                if let Some(entity) = extract_balanced(text, i) {
                    let end = i + entity.len();
                    results.push(entity);
                    i = end;
                    continue;
                }
                // Unbalanced — skip this opener
            }
            _ => {}
        }
        i += 1;
    }

    results
}

/// Starting from `start` (which must be `{` or `[`), find the matching close.
fn extract_balanced(text: &str, start: usize) -> Option<&str> {
    let bytes = text.as_bytes();
    let len = bytes.len();
    let open = bytes[start];
    let close = if open == b'{' { b'}' } else { b']' };

    let mut depth: usize = 0;
    let mut in_string = false;
    let mut escape_next = false;
    let mut i = start;

    while i < len {
        let b = bytes[i];

        if escape_next {
            escape_next = false;
            i += 1;
            continue;
        }

        if in_string {
            match b {
                b'\\' => escape_next = true,
                b'"' => in_string = false,
                _ => {}
            }
            i += 1;
            continue;
        }

        match b {
            b'"' => in_string = true,
            b'{' | b'[' => depth += 1,
            b'}' | b']' => {
                depth -= 1;
                if depth == 0 && b == close {
                    return Some(&text[start..=i]);
                }
            }
            _ => {}
        }

        i += 1;
    }

    // No balanced close found — return from start to end as best-effort
    // (sanitize_json will attempt to close it)
    None
}

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

    #[test]
    fn bare_object() {
        assert_eq!(extract_json_entity(r#"{"a": 1}"#), Some(r#"{"a": 1}"#));
    }

    #[test]
    fn bare_array() {
        assert_eq!(extract_json_entity("[1, 2, 3]"), Some("[1, 2, 3]"));
    }

    #[test]
    fn object_with_prose_before() {
        let text = r#"Here's the data: {"a": 1}"#;
        assert_eq!(extract_json_entity(text), Some(r#"{"a": 1}"#));
    }

    #[test]
    fn object_with_prose_after() {
        let text = "{\"a\": 1}\n\nExplanation follows.";
        assert_eq!(extract_json_entity(text), Some("{\"a\": 1}"));
    }

    #[test]
    fn object_with_prose_both() {
        let text = "Sure! {\"a\": 1} Hope that helps!";
        assert_eq!(extract_json_entity(text), Some("{\"a\": 1}"));
    }

    #[test]
    fn nested_braces() {
        let text = r#"{"outer": {"inner": [1, 2]}}"#;
        assert_eq!(extract_json_entity(text), Some(text));
    }

    #[test]
    fn braces_in_string() {
        let text = r#"{"msg": "use {x} here"}"#;
        assert_eq!(extract_json_entity(text), Some(text));
    }

    #[test]
    fn escaped_quote_in_string() {
        let text = r#"{"msg": "say \"hello\""}"#;
        assert_eq!(extract_json_entity(text), Some(text));
    }

    #[test]
    fn multiple_entities_returns_first() {
        let text = r#"{"a": 1} some text {"b": 2}"#;
        assert_eq!(extract_json_entity(text), Some(r#"{"a": 1}"#));
    }

    #[test]
    fn empty_object() {
        assert_eq!(extract_json_entity("{}"), Some("{}"));
    }

    #[test]
    fn empty_array() {
        assert_eq!(extract_json_entity("[]"), Some("[]"));
    }

    #[test]
    fn no_json() {
        assert_eq!(extract_json_entity("No data available"), None);
    }

    #[test]
    fn unclosed_bracket() {
        assert_eq!(extract_json_entity("{\"a\": 1"), None);
    }

    #[test]
    fn commas_in_strings_preserved() {
        let text = r#"{"msg": "hello, world"}"#;
        assert_eq!(extract_json_entity(text), Some(text));
    }

    // extract_all_json_entities tests

    #[test]
    fn all_entities_multiple() {
        let text = r#"{"a": 1} some text {"b": 2} more {"c": 3}"#;
        let entities = extract_all_json_entities(text);
        assert_eq!(entities.len(), 3);
        assert_eq!(entities[0], r#"{"a": 1}"#);
        assert_eq!(entities[1], r#"{"b": 2}"#);
        assert_eq!(entities[2], r#"{"c": 3}"#);
    }

    #[test]
    fn all_entities_mixed_types() {
        let text = r#"{"a": 1} and [1, 2, 3]"#;
        let entities = extract_all_json_entities(text);
        assert_eq!(entities.len(), 2);
        assert_eq!(entities[0], r#"{"a": 1}"#);
        assert_eq!(entities[1], "[1, 2, 3]");
    }

    #[test]
    fn all_entities_single() {
        let text = r#"{"a": 1}"#;
        let entities = extract_all_json_entities(text);
        assert_eq!(entities.len(), 1);
    }

    #[test]
    fn all_entities_none() {
        let entities = extract_all_json_entities("no json here");
        assert!(entities.is_empty());
    }
}
