//! Functions for parsing the output of LLMs, including YAML and Markdown.
//!
//! This module provides a set of functions that allow you to parse and extract useful information from the output of Large Language Models (LLMs) in YAML and Markdown formats. These functions can be used to transform the LLM output into a more structured and usable format, enabling seamless integration with your applications.
//!
//! Key features include:
//! - Parsing YAML and Markdown content produced by LLMs
//! - Handling common edge cases and being lenient with LLM outputs
//! - Extracting and deserializing YAML objects from text
//!
//! With these functions, you can easily work with the outputs of LLMs, simplifying the process of integrating LLMs into your applications and workflows.

use markdown::{
    mdast::{Code, Node},
    to_mdast, ParseOptions,
};
use serde::de::DeserializeOwned;
use serde_yaml::Value;
use thiserror::Error;

/// Errors occuring when parsing
#[derive(Error, Debug)]
#[error(transparent)]
pub struct ExtractionError(#[from] ExtractionErrorImpl);

/// Enum representing the different error cases that can occur during YAML extraction.
#[derive(Error, Debug)]
enum ExtractionErrorImpl {
    /// The YAML content was valid, but it did not match the expected format.
    #[error("The YAML was valid, but it didn't match the expected format: {0}")]
    YamlFoundButFormatWrong(serde_yaml::Error),

    /// An error occurred while parsing the YAML content.
    #[error("YAML parsing failed with: {0}")]
    ParseError(#[from] serde_yaml::Error),

    /// No YAML content was found to parse.
    #[error("The string to parse was empty")]
    NoneFound,
}

impl ExtractionErrorImpl {
    /// Determines the most representative error between two instances of `ExtractionErrorImpl`.
    ///
    /// The function prefers `YamlFoundButFormatWrong` errors over `ParseError` errors,
    /// and `ParseError` errors over `NoneFound` errors.
    fn most_representative(a: Self, b: Self) -> Self {
        match (&a, &b) {
            (Self::YamlFoundButFormatWrong(_), _) => a,
            (_, Self::YamlFoundButFormatWrong(_)) => b,
            (Self::ParseError(_), _) => a,
            (_, Self::ParseError(_)) => b,
            _ => a,
        }
    }
}

/// Attempts to extract YAML content from a given code block and deserialize it into the specified type.
///
/// # Arguments
///
/// * `code_block` - A string slice containing the YAML content to be extracted and deserialized.
///
/// # Returns
///
/// * `Ok(T)` - If the YAML content is successfully extracted and deserialized into the specified type.
/// * `Err(ExtractionErrorImpl)` - If an error occurs during extraction or deserialization.
///
/// # Type Parameters
///
/// * `T: DeserializeOwned` - The type into which the extracted YAML content should be deserialized.
fn extract_yaml<T: DeserializeOwned>(code_block: &str) -> Result<T, ExtractionErrorImpl> {
    // Ensure that the input code block is not empty.
    let code_block = Some(code_block)
        .filter(|s| !s.is_empty())
        .ok_or_else(|| ExtractionErrorImpl::NoneFound)?;

    // Parse the code block as YAML.
    let yaml: Value = serde_yaml::from_str(code_block)?;

    // Attempt to deserialize the YAML into the specified type, handling any format issues.
    serde_yaml::from_value(yaml).map_err(ExtractionErrorImpl::YamlFoundButFormatWrong)
}

/// Attempts to find a YAML object in a string and deserialize it into the specified type.
/// # Arguments
///
/// * `text` - A string slice containing the document which may contain YAML content.
///
/// # Returns
///
/// * `Ok(Vec<T>)` - If the YAML content is successfully extracted and deserialized into the specified type.
/// * `Err(ExtractionError)` - If an error occurs during extraction or deserialization.
///
/// # Type Parameters
///
/// * `T: DeserializeOwned` - The type into which the extracted YAML content should be deserialized.
///
/// # Examples
///
/// It handles the obvious case where it is just YAML.
///
/// ```
/// #[derive(serde::Deserialize)]
/// struct Dummy {
///    hello: String
/// }
/// use llm_chain::parsing::find_yaml;
/// let data = "
/// hello: world
/// ";
/// let data: Vec<Dummy> = find_yaml(data).unwrap();
/// assert_eq!(data[0].hello, "world");
/// ```
///
/// It handles the case where it is in a code block.
///
/// ```
/// use llm_chain::parsing::find_yaml;
/// // NOTE: we are escaping the backticks because this is a doc test.
/// let data = "
/// \u{60}``yaml
/// hello: world
/// \u{60}``
/// ";
/// find_yaml::<serde_yaml::Value>(data).unwrap();
/// ```
pub fn find_yaml<T: DeserializeOwned>(text: &str) -> Result<Vec<T>, ExtractionError> {
    let mut current_error = ExtractionErrorImpl::NoneFound;
    if text.is_empty() {
        return Err(current_error.into());
    }

    // Attempt YAML parsing if it doesn't look like markdown output.
    if !text.starts_with("```") {
        match extract_yaml(text) {
            Ok(o) => return Ok(vec![o]),
            Err(e) => current_error = ExtractionErrorImpl::most_representative(current_error, e),
        }
    }

    // Parse the input text as markdown.
    let options = ParseOptions::default();
    let ast = to_mdast(text, &options).expect("we're not using MDX, so this should never fail");

    // Nodes to visit.
    let mut nodes = vec![ast];

    let mut found: Vec<T> = Vec::new();
    while let Some(node) = nodes.pop() {
        if let Some(children) = node.children() {
            children.iter().for_each(|child| nodes.push(child.clone()));
        }

        // Check for code blocks containing YAML.
        if let Node::Code(Code { value, lang, .. }) = node {
            let lang = lang.unwrap_or_default();
            match lang.as_str() {
                "yaml" | "yml" | "json" | "" => {
                    let code_block = value.as_str();
                    match extract_yaml(code_block) {
                        Ok(o) => found.push(o),
                        Err(e) => {
                            current_error =
                                ExtractionErrorImpl::most_representative(current_error, e)
                        }
                    }
                }
                _ => {}
            }
        }
    }
    if !found.is_empty() {
        Ok(found)
    } else {
        Err(current_error.into())
    }
}