lc-core 0.16.0

use async_trait::async_trait;
use futures_util::Stream;
use serde::de::DeserializeOwned;
use std::collections::HashMap;
use std::marker::PhantomData;
use std::pin::Pin;

use super::base::{BaseOutputParser, OutputParserError, OutputParserResult};
use crate::language_models::LLMResult;
use crate::runnables::{Runnable, RunnableConfig};

/// 结构化输出解析器
///
/// 将 LLM 输出的键值对格式（每行一个 `key: value`）解析为 HashMap。
/// 适用于 LLM 以非 JSON 格式输出结构化信息的场景。
///
/// # 格式
/// 输入格式应为每行一个 `key: value`，例如：
/// ```text
/// 姓名: 张三
/// 年龄: 28
/// 城市: 北京
/// ```
///
/// # 示例
/// ```ignore
/// use langchainrust::output_parsers::StructuredOutputParser;
///
/// let parser = StructuredOutputParser::new();
/// let result = parser.parse("姓名: 张三\n年龄: 28").await?;
/// assert_eq!(result.get("姓名").unwrap(), "张三");
/// ```
pub struct StructuredOutputParser {
    /// 键值对之间的分隔符
    separator: char,
}

impl StructuredOutputParser {
    /// 创建使用默认分隔符(`:`)的结构化输出解析器。
    pub fn new() -> Self {
        Self { separator: ':' }
    }

    /// 使用自定义分隔符创建解析器
    pub fn with_separator(separator: char) -> Self {
        Self { separator }
    }
}

impl Default for StructuredOutputParser {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl BaseOutputParser<HashMap<String, String>> for StructuredOutputParser {
    async fn parse(&self, text: &str) -> OutputParserResult<HashMap<String, String>> {
        let mut map = HashMap::new();

        for line in text.lines() {
            let line = line.trim();
            if line.is_empty() {
                continue;
            }

            if let Some(pos) = line.find(self.separator) {
                // `pos` 是分隔符首字节的字节索引;多字节分隔符(如全角 `：`)时
                // pos+1 会落在字符内部导致切片 panic,须按分隔符 UTF-8 宽度跳过
                let sep_len = self.separator.len_utf8();
                let key = line[..pos].trim().to_string();
                let value = line[pos + sep_len..].trim().to_string();

                if !key.is_empty() {
                    map.insert(key, value);
                }
            }
        }

        Ok(map)
    }

    fn get_format_instructions(&self) -> String {
        format!(
            "请按以下格式输出（每行一个键值对，使用 '{}' 分隔）：\n键{}值",
            self.separator, self.separator
        )
    }
}

#[async_trait]
impl Runnable<LLMResult, HashMap<String, String>> for StructuredOutputParser {
    type Error = OutputParserError;

    async fn invoke(
        &self,
        input: LLMResult,
        _config: Option<RunnableConfig>,
    ) -> Result<HashMap<String, String>, Self::Error> {
        self.parse(&input.content).await
    }

    async fn stream(
        &self,
        input: LLMResult,
        _config: Option<RunnableConfig>,
    ) -> Result<
        Pin<Box<dyn Stream<Item = Result<HashMap<String, String>, Self::Error>> + Send>>,
        Self::Error,
    > {
        let result = self.parse(&input.content).await?;
        let stream = futures_util::stream::once(async move { Ok(result) });
        Ok(Box::pin(stream))
    }
}

/// 类型化输出解析器
///
/// 将 LLM 输出的 JSON 字符串解析为指定的 Rust 结构体。
/// 相当于 Python LangChain 的 `PydanticOutputParser`（使用 serde 替代 pydantic）。
///
/// 需要目标类型实现 `serde::Deserialize`。
///
/// # 示例
/// ```ignore
/// use serde::Deserialize;
/// use langchainrust::output_parsers::TypedOutputParser;
///
/// #[derive(Deserialize, Debug, PartialEq)]
/// struct Person {
///     name: String,
///     age: u32,
/// }
///
/// let parser = TypedOutputParser::<Person>::new();
/// let person = parser.parse(r#"{"name": "Alice", "age": 30}"#).await?;
/// assert_eq!(person.name, "Alice");
/// ```
pub struct TypedOutputParser<T> {
    _phantom: PhantomData<T>,
}

impl<T> TypedOutputParser<T> {
    /// 创建类型化输出解析器。
    pub fn new() -> Self {
        Self {
            _phantom: PhantomData,
        }
    }
}

impl<T: DeserializeOwned> Default for TypedOutputParser<T> {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl<T: DeserializeOwned + Send + Sync + 'static> BaseOutputParser<T> for TypedOutputParser<T> {
    async fn parse(&self, text: &str) -> OutputParserResult<T> {
        let text = text.trim();

        // 尝试从 Markdown 代码块中提取 JSON
        let json_str = Self::extract_from_markdown(text).unwrap_or(text);

        // 先尝试解析为 Value 验证合法性
        serde_json::from_str::<serde_json::Value>(json_str)
            .map_err(|e| OutputParserError::JsonError(format!("input is not valid JSON: {}", e)))?;

        // 反序列化为目标类型
        serde_json::from_str::<T>(json_str).map_err(|e| {
            OutputParserError::TypeError(format!(
                "type deserialization failed (check whether the JSON fields match): {}",
                e
            ))
        })
    }
}

impl<T: DeserializeOwned> TypedOutputParser<T> {
    /// 从 Markdown 代码块中提取 JSON 字符串
    fn extract_from_markdown(text: &str) -> Option<&str> {
        // 尝试 ```json ... ```
        if let Some(start) = text.find("```json") {
            let after = &text[start + 7..];
            if let Some(end) = after.find("```") {
                return Some(after[..end].trim());
            }
        }
        // 尝试 ``` ... ```
        if let Some(start) = text.find("```") {
            let after = &text[start + 3..];
            let after = after.trim();
            let skip = after.find('\n').unwrap_or(0);
            let after = &after[skip..].trim();
            if let Some(end) = after.find("```") {
                return Some(after[..end].trim());
            }
        }
        None
    }
}

#[async_trait]
impl<T: DeserializeOwned + Send + Sync + 'static> Runnable<LLMResult, T> for TypedOutputParser<T> {
    type Error = OutputParserError;

    async fn invoke(
        &self,
        input: LLMResult,
        _config: Option<RunnableConfig>,
    ) -> Result<T, Self::Error> {
        self.parse(&input.content).await
    }

    async fn stream(
        &self,
        input: LLMResult,
        _config: Option<RunnableConfig>,
    ) -> Result<Pin<Box<dyn Stream<Item = Result<T, Self::Error>> + Send>>, Self::Error> {
        let result = self.parse(&input.content).await?;
        let stream = futures_util::stream::once(async move { Ok(result) });
        Ok(Box::pin(stream))
    }
}

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

    #[tokio::test]
    async fn test_structured_parser_default_colon() {
        let parser = StructuredOutputParser::new();
        let map = parser.parse("姓名: 张三\n年龄: 28").await.unwrap();
        assert_eq!(map.get("姓名").unwrap(), "张三");
        assert_eq!(map.get("年龄").unwrap(), "28");
    }

    #[tokio::test]
    async fn test_structured_parser_fullwidth_separator() {
        // 全角冒号 3 字节:修复前按 pos+1 切片会切在字符内部 panic
        let parser = StructuredOutputParser::with_separator('：');
        let map = parser.parse("姓名：张三\n年龄：28").await.unwrap();
        assert_eq!(map.get("姓名").unwrap(), "张三");
        assert_eq!(map.get("年龄").unwrap(), "28");
    }

    #[tokio::test]
    async fn test_structured_parser_runnable_invoke() {
        let parser = StructuredOutputParser::new();
        let map = parser
            .invoke(
                LLMResult {
                    content: "状态: 成功".to_string(),
                    ..Default::default()
                },
                None,
            )
            .await
            .unwrap();
        assert_eq!(map.get("状态").unwrap(), "成功");
    }
}