lellm-provider 0.1.2

//! ToolCall Accumulator — 按 index 聚合增量 delta。
//!
//! 以 index 为 key，因为很多 Provider 的第一批 delta 只有 index 而没有 id。
//!
//! **Structured Output — 组合拳策略，最大化 JSON 合规率。**
//!
//! 策略链（按优先级）：
//! 1. **Tool Use 模式**（默认）— 定义 tool + tool_choice 强制调用，99.8% 合规率
//! 2. **6 层兜底解析** — 救回 90% 的小翻车（截断、尾逗号、单引号、markdown 代码块等）

use lellm_core::{LlmError, ToolCall};

/// 工具调用增量 — 统一格式，吸收所有 Provider 差异
#[derive(Debug)]
pub struct ToolCallDelta {
    /// 工具调用在消息中的位置索引（用于聚合）
    pub index: usize,
    /// 工具调用 ID（可能延迟出现）
    pub id: Option<String>,
    /// 工具名称（可能延迟出现）
    pub name: Option<String>,
    /// 参数增量片段（最终拼接为完整 JSON）
    pub arguments_delta: Option<String>,
}

/// ToolCall 增量组装器 — 按 index 聚合。
///
/// 独立状态机，可单独测试。
pub struct ToolCallAccumulator {
    current: std::collections::HashMap<usize, PendingToolCall>,
}

struct PendingToolCall {
    id: Option<String>,
    name: Option<String>,
    arguments: String,
}

impl ToolCallAccumulator {
    pub fn new() -> Self {
        Self {
            current: std::collections::HashMap::new(),
        }
    }

    /// 接收统一的 ToolCallDelta 增量并组装
    pub fn push(&mut self, delta: &ToolCallDelta) {
        let entry = self
            .current
            .entry(delta.index)
            .or_insert_with(|| PendingToolCall {
                id: None,
                name: None,
                arguments: String::new(),
            });

        if let Some(ref id) = delta.id {
            entry.id = Some(id.clone());
        }
        if let Some(ref name) = delta.name {
            entry.name = Some(name.clone());
        }
        if let Some(ref d) = delta.arguments_delta {
            entry.arguments.push_str(d);
        }
    }

    /// 完成组装，返回完整的 ToolCall 列表（按 index 排序）
    pub fn finalize(self) -> Result<Vec<ToolCall>, LlmError> {
        let mut entries: Vec<_> = self.current.into_iter().collect();
        entries.sort_by_key(|&(idx, _)| idx);

        let mut result = Vec::new();
        for (_index, pending) in entries {
            let id = pending.id.unwrap_or_else(|| "unknown".to_string());
            let name = pending.name.unwrap_or_else(|| "unknown".to_string());
            let arguments: serde_json::Value = if pending.arguments.is_empty() {
                serde_json::Value::Object(Default::default())
            } else {
                robust_parse(&pending.arguments)
            };
            result.push(ToolCall {
                id,
                name,
                arguments,
            });
        }
        Ok(result)
    }
}

// ─── 6 层兜底解析 ──────────────────────────────────────────────

/// 6 层递进解析，最大化 JSON 合规率。
///
/// 1. 直接解析
/// 2. 剥离 markdown 代码块
/// 3. 找到第一个 { 和最后一个 }
/// 4. 修复常见错误（尾逗号、单引号）
/// 5. 再试
/// 6. 返回原始字符串（让下游 ToolError::InvalidInput 处理）
fn robust_parse(text: &str) -> serde_json::Value {
    let trimmed = text.trim();

    // Empty input: return empty object (finalize() handles this path, but be defensive)
    if trimmed.is_empty() {
        return serde_json::Value::Object(Default::default());
    }

    // Layer 1: Direct parse
    if let Ok(result) = serde_json::from_str::<serde_json::Value>(&trimmed) {
        return result;
    }

    // Layer 2: Strip markdown code blocks
    let stripped = strip_codeblocks(&trimmed);
    if let Ok(result) = serde_json::from_str::<serde_json::Value>(&stripped) {
        tracing::debug!(
            original_preview = %trimmed.chars().take(80).collect::<String>(),
            "structured output: stripped markdown codeblocks"
        );
        return result;
    }

    // Layer 3: Extract outermost { ... }
    if let Some(json_str) = extract_braces(&stripped) {
        if let Ok(result) = serde_json::from_str::<serde_json::Value>(&json_str) {
            tracing::debug!(
                original_preview = %trimmed.chars().take(80).collect::<String>(),
                "structured output: extracted braces"
            );
            return result;
        }
    }

    // Layer 4: Fix common errors — trailing commas, single quotes
    let fixed = fix_common_errors(&stripped);
    if let Some(json_str) = extract_braces(&fixed) {
        if let Ok(result) = serde_json::from_str::<serde_json::Value>(&json_str) {
            tracing::debug!(
                original_preview = %trimmed.chars().take(80).collect::<String>(),
                "structured output: fixed common json errors"
            );
            return result;
        }
    }

    // Layer 5: Try as Value via round-trip (normalize whitespace, etc.)
    if let Ok(v) = serde_json::from_str::<serde_json::Value>(&fixed) {
        return v;
    }

    // Layer 6: Return raw string — downstream will get ToolError::InvalidInput
    tracing::warn!(
        raw_preview = %trimmed.chars().take(120).collect::<String>(),
        "structured output: all parse layers failed, returning raw string"
    );
    serde_json::Value::String(trimmed.to_string())
}

/// 剥离 markdown 代码块标记
fn strip_codeblocks(text: &str) -> String {
    let mut result = text.to_string();
    result = result.replace("```json\n", "").replace("```json", "");
    result = result.replace("```\n", "").replace("```", "");
    result.trim().to_string()
}

/// 修复常见 JSON 错误
fn fix_common_errors(text: &str) -> String {
    let mut s = text.to_string();
    // Fix trailing commas before } or ]
    s = s
        .replace(", }", "}")
        .replace(",\t}", "}")
        .replace(",}", "}");
    s = s
        .replace(", ]", "]")
        .replace(",\t]", "]")
        .replace(",]", "]");
    // Fix single quotes to double quotes using a lightweight state machine
    // that respects double-quoted strings (don't touch apostrophes inside them).
    s = normalize_single_quote_json(&s);
    s
}

/// Normalize single-quoted JSON to double-quoted JSON using a lightweight state machine.
///
/// Respects double-quoted strings so apostrophes inside them are preserved:
/// `{'text': "It's"}` → `{"text": "It's"}`
///
/// Also handles escape sequences: `{"a": "test\"ed"}` — escaped quotes don't flip state.
fn normalize_single_quote_json(s: &str) -> String {
    let mut out = String::new();
    let mut in_double = false;
    let mut in_single = false;
    let mut escaped = false;

    for ch in s.chars() {
        match ch {
            '"' if !escaped && !in_single => {
                in_double = !in_double;
                out.push(ch);
            }
            '\'' if !escaped && !in_double => {
                in_single = !in_single;
                out.push('"');
            }
            _ => out.push(ch),
        }
        escaped = ch == '\\' && !escaped;
    }

    out
}

/// Extract the outermost { ... } from the content.
fn extract_braces(content: &str) -> Option<String> {
    let start = content.find('{')?;
    let mut depth = 0i32;
    let mut end = None;
    let mut in_string = false;
    let mut escaped = false;

    for (i, c) in content[start..].char_indices() {
        if in_string {
            if escaped {
                escaped = false;
            } else if c == '\\' {
                escaped = true;
            } else if c == '"' {
                in_string = false;
            }
        } else {
            match c {
                '"' => in_string = true,
                '{' => depth += 1,
                '}' => {
                    depth -= 1;
                    if depth == 0 {
                        end = Some(start + i + 1);
                        break;
                    }
                }
                _ => {}
            }
        }
    }
    end.map(|e| content[start..e].to_string())
}

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

    #[test]
    fn accumulates_delta_into_tool_call() {
        let mut acc = ToolCallAccumulator::new();

        // 第一批 delta: 只有 index + name
        acc.push(&ToolCallDelta {
            index: 0,
            id: None,
            name: Some("read_file".into()),
            arguments_delta: None,
        });

        // 第二批 delta: arguments 片段
        acc.push(&ToolCallDelta {
            index: 0,
            id: Some("tc_123".into()),
            name: None,
            arguments_delta: Some(r#""{"path""#.into()),
        });

        // 第三批 delta: arguments 片段
        acc.push(&ToolCallDelta {
            index: 0,
            id: None,
            name: None,
            arguments_delta: Some(r#": "/test"}}"#.into()),
        });

        let calls = acc.finalize().unwrap();
        assert_eq!(calls.len(), 1);
        assert_eq!(calls[0].id, "tc_123");
        assert_eq!(calls[0].name, "read_file");
    }

    #[test]
    fn empty_accumulator_returns_empty() {
        let acc = ToolCallAccumulator::new();
        let calls = acc.finalize().unwrap();
        assert!(calls.is_empty());
    }

    // ─── robust_parse 测试 ───

    #[test]
    fn robust_parse_direct() {
        let result = robust_parse(r#"{"action":"deploy","job_name":"ds-pkg"}"#);
        assert!(result.is_object());
        assert_eq!(result["action"], "deploy");
    }

    #[test]
    fn robust_parse_codeblock() {
        let result = robust_parse("```json\n{\"action\":\"build\"}\n```");
        assert_eq!(result["action"], "build");
    }

    #[test]
    fn robust_parse_braces() {
        let result = robust_parse("Here is the result: {\"action\":\"query\"} done.");
        assert_eq!(result["action"], "query");
    }

    #[test]
    fn robust_parse_trailing_comma() {
        let result = robust_parse(r#"{"action":"deploy","job_name":"ds-pkg","branch":"main",}"#);
        assert_eq!(result["action"], "deploy");
    }

    #[test]
    fn robust_parse_single_quotes() {
        let result = robust_parse(r#"{'action':'deploy','job_name':'ds-pkg'}"#);
        assert_eq!(result["action"], "deploy");
    }

    #[test]
    fn robust_parse_nested_braces() {
        let result = robust_parse(r#"{"config":{"nested":true}}"#);
        assert_eq!(result["config"]["nested"], true);
    }

    #[test]
    fn robust_parse_braces_in_string() {
        let result = robust_parse(r#"{"text":"hello {world}"}"#);
        assert_eq!(result["text"], "hello {world}");
    }

    #[test]
    fn robust_parse_empty_returns_object() {
        let result = robust_parse("");
        assert!(result.is_object());
    }

    #[test]
    fn strip_codeblocks_removes_markers() {
        assert_eq!(strip_codeblocks("```json\n{\"a\":1}\n```"), r#"{"a":1}"#);
        assert_eq!(strip_codeblocks("```\n{\"a\":1}\n```"), r#"{"a":1}"#);
    }

    #[test]
    fn fix_common_errors_removes_trailing_commas() {
        assert_eq!(fix_common_errors(r#"{"a":1,}"#), r#"{"a":1}"#);
        assert_eq!(fix_common_errors(r#"{"a":[1,]}"#), r#"{"a":[1]}"#);
    }

    #[test]
    fn normalize_single_quote_json_pure_single_quotes() {
        assert_eq!(
            normalize_single_quote_json(r#"{'action':'deploy'}"#),
            r#"{"action":"deploy"}"#
        );
    }

    #[test]
    fn normalize_single_quote_json_preserves_apostrophe_in_double_string() {
        assert_eq!(
            normalize_single_quote_json(r#"{'text': "It's"}"#),
            r#"{"text": "It's"}"#
        );
    }

    #[test]
    fn normalize_single_quote_json_no_change_if_double() {
        assert_eq!(
            normalize_single_quote_json(r#"{"a": "b"}"#),
            r#"{"a": "b"}"#
        );
    }

    #[test]
    fn normalize_single_quote_json_respects_escape() {
        assert_eq!(
            normalize_single_quote_json(r#"{"a": "test\"ed"}"#),
            r#"{"a": "test\"ed"}"#
        );
    }

    #[test]
    fn normalize_single_quote_json_double_backslash() {
        // \\ should reset escaped state, so next " flips in_double
        assert_eq!(
            normalize_single_quote_json(r#"{"a": "C:\\\"test"}"#),
            r#"{"a": "C:\\\"test"}"#
        );
    }

    #[test]
    fn robust_parse_mixed_quotes_with_trailing_comma() {
        // Single quotes for keys, double-quoted value with apostrophe, trailing comma
        let result = robust_parse(r#"{'text': "It's a test", 'count': 5,}"#);
        assert_eq!(result["text"], "It's a test");
        assert_eq!(result["count"], 5);
    }

    #[test]
    fn extract_braces_finds_outermost() {
        assert_eq!(
            extract_braces("hello {\"a\":1} world"),
            Some(r#"{"a":1}"#.into())
        );
        assert_eq!(
            extract_braces("{\"a\":{\"b\":1}}"),
            Some(r#"{"a":{"b":1}}"#.into())
        );
        assert_eq!(
            extract_braces(r#"{"text":"hello {world}"}"#),
            Some(r#"{"text":"hello {world}"}"#.into())
        );
    }
}