use crate::error::{CognisError, Result};
use serde_json::Value;
pub fn parse_partial_json(s: &str) -> Result<Value> {
if let Ok(v) = serde_json::from_str(s) {
return Ok(v);
}
let trimmed = s.trim();
let mut attempt = trimmed.to_string();
let open_braces = attempt.chars().filter(|c| *c == '{').count();
let close_braces = attempt.chars().filter(|c| *c == '}').count();
let open_brackets = attempt.chars().filter(|c| *c == '[').count();
let close_brackets = attempt.chars().filter(|c| *c == ']').count();
if attempt.ends_with(',') {
attempt.pop();
}
for _ in 0..(open_brackets.saturating_sub(close_brackets)) {
attempt.push(']');
}
for _ in 0..(open_braces.saturating_sub(close_braces)) {
attempt.push('}');
}
serde_json::from_str(&attempt)
.map_err(|e| CognisError::Other(format!("Failed to parse partial JSON: {}", e)))
}
pub fn parse_json_markdown(text: &str) -> Result<Value> {
let trimmed = text.trim();
let json_str = if trimmed.starts_with("```") {
let after = trimmed
.strip_prefix("```json")
.or_else(|| trimmed.strip_prefix("```JSON"))
.or_else(|| trimmed.strip_prefix("```"))
.unwrap_or(trimmed);
after.trim().strip_suffix("```").unwrap_or(after).trim()
} else {
trimmed
};
serde_json::from_str(json_str)
.map_err(|e| CognisError::Other(format!("Failed to parse JSON from markdown: {}", e)))
}
pub fn parse_and_check_json_markdown(text: &str, expected_keys: &[&str]) -> Result<Value> {
let parsed = parse_json_markdown(text)?;
if let Value::Object(ref map) = parsed {
for key in expected_keys {
if !map.contains_key(*key) {
return Err(CognisError::Other(format!(
"Missing expected key '{}' in JSON output",
key
)));
}
}
}
Ok(parsed)
}