unai 0.0.1

use schemars::JsonSchema;
use serde::de::DeserializeOwned;

use super::completion::Completion;
use super::content::Content;
use super::llm::{GenerateCompletionError, GenerateCompletionInput, LLM};
use super::prompt::{Message, Prompt};
use super::settings::Settings;

#[derive(Debug, Clone, Default)]
pub struct GenerateTextInput {
    /// The prompt to generate a completion for
    pub prompt: Prompt,

    /// The settings to use for the generation
    pub settings: Option<Settings>,
}

#[derive(Debug, Clone)]
pub struct GenerateTextOutput {
    /// The generated text
    pub text: String,

    /// The full completion
    pub completion: Completion,
}

/// Generate text
pub async fn generate_text(
    llm: &dyn LLM,
    input: GenerateTextInput,
) -> Result<GenerateTextOutput, GenerateCompletionError> {
    let generate_completion_input = GenerateCompletionInput {
        prompt: input.prompt,
        ..Default::default()
    };

    let completion = llm.generate_completion(generate_completion_input).await?;

    let mut completion_text = String::new();

    for content in completion.content.iter() {
        if let Content::Text(text) = content {
            completion_text.push_str(text);
        }
    }

    Ok(GenerateTextOutput {
        text: completion_text,
        completion,
    })
}

#[derive(Debug, Clone, Default)]
pub struct GenerateJsonInput {
    /// The prompt to generate a completion for
    pub prompt: Prompt,

    /// The settings to use for the generation
    pub settings: Option<Settings>,

    /// JSON schema for structured generation
    pub schema: serde_json::Value,
}

#[derive(Debug, Clone)]
pub struct GenerateJsonOutput {
    /// The generated JSON value
    pub json: serde_json::Value,

    /// The full completion
    pub completion: Completion,
}

#[derive(Debug, thiserror::Error)]
pub enum GenerateJsonError {
    #[error("deserialization error: {0}")]
    Deserialization(#[from] serde_json::Error),

    #[error("deserialization error: {0}")]
    InvalidJson(#[from] Box<jsonschema::ValidationError<'static>>),

    #[error("an error occurred while generating the completion: {0}")]
    GenerateCompletion(#[from] GenerateCompletionError),
}

pub type GenerateJsonResult = Result<GenerateJsonOutput, GenerateJsonError>;

/// Generate structured JSON output based on the provided JSON schema.
///
/// The result from the completion is validated against the provided schema.
pub async fn generate_json(llm: &dyn LLM, input: GenerateJsonInput) -> GenerateJsonResult {
    let generate_completion_input = GenerateCompletionInput {
        prompt: input.prompt,
        settings: input.settings,
        schema: Some(input.schema.clone()),
    };

    let completion = llm.generate_completion(generate_completion_input).await?;

    let mut completion_str = String::new();

    for content in completion.content.iter() {
        if let Content::Text(text) = content {
            completion_str.push_str(text);
        }
    }

    let json: serde_json::Value = serde_json::from_str(&completion_str)?;

    if let Some(err) = jsonschema::validate(&input.schema, &json).err() {
        return Err(GenerateJsonError::InvalidJson(Box::new(err.to_owned())));
    };

    Ok(GenerateJsonOutput { json, completion })
}

#[derive(Debug, Clone, Default)]
pub struct GenerateValueInput {
    /// The prompt to generate a completion for
    pub prompt: Prompt,

    /// The settings to use for the generation
    pub settings: Option<Settings>,
}

impl From<String> for GenerateValueInput {
    fn from(message: String) -> Self {
        Self {
            prompt: Prompt::from(message),
            ..Default::default()
        }
    }
}

impl From<&str> for GenerateValueInput {
    fn from(message: &str) -> Self {
        Self::from(message.to_owned())
    }
}

impl From<Prompt> for GenerateValueInput {
    fn from(prompt: Prompt) -> Self {
        Self {
            prompt,
            ..Default::default()
        }
    }
}

impl From<Message> for GenerateValueInput {
    fn from(message: Message) -> Self {
        Self::from(Prompt::from(message))
    }
}

#[derive(Debug, Clone)]
pub struct GenerateValueOutput<T> {
    /// The generated value
    pub value: T,

    /// The full completion
    pub completion: Completion,
}

pub type GenerateValueResult<T> = Result<GenerateValueOutput<T>, GenerateJsonError>;

/// Generate a value from a type
///
/// # Type Parameters
/// - `T`: The type to generate. Must implement `schemars::JsonSchema` and `serde::DeserializeOwned`.
pub async fn generate_value<T: JsonSchema + DeserializeOwned>(
    llm: &dyn LLM,
    input: GenerateValueInput,
) -> GenerateValueResult<T> {
    let json_schema = schemars::schema_for!(T);
    let json_schema_value = json_schema.to_value();

    let generate_json_input = GenerateJsonInput {
        prompt: input.prompt,
        settings: input.settings,
        schema: json_schema_value,
    };

    let generate_json_output = generate_json(llm, generate_json_input).await?;
    let value: T = serde_json::from_value(generate_json_output.json)?;

    Ok(GenerateValueOutput {
        value,
        completion: generate_json_output.completion,
    })
}