//! `MetaCrate` message/tool types and the configured Mentra provider.

#![allow(clippy::missing_errors_doc)]

use crate::config::LlmConnection;
use crate::types::{
    BoundedText, BoundedVec, MAX_BODY_BYTES, MAX_IDENTIFIER_BYTES, MAX_MESSAGE_BYTES,
    MAX_TOOL_CALLS, MessageRole, ProposedToolCall,
};
use serde_json::{Map, Value, json};
use std::collections::{BTreeMap, BTreeSet};
use std::error::Error;
use std::fmt;

const MAX_SCHEMA_PROPERTIES: usize = 128;
const MAX_SCHEMA_DEPTH: usize = 16;
const MAX_SCHEMA_NODES: usize = 1_024;

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum ContentPart {
    Text(BoundedText<MAX_MESSAGE_BYTES>),
    Image {
        url: BoundedText<MAX_BODY_BYTES>,
        detail: ImageDetail,
    },
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ImageDetail {
    Auto,
    Low,
    High,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct CompletionMessage {
    pub role: MessageRole,
    pub content: BoundedVec<ContentPart, 16>,
    pub tool_call_id: Option<BoundedText<MAX_IDENTIFIER_BYTES>>,
    pub proposed_calls: BoundedVec<ProposedToolCall, MAX_TOOL_CALLS>,
}

impl CompletionMessage {
    pub fn text(role: MessageRole, text: impl Into<String>) -> Result<Self, LlmError> {
        let mut content = BoundedVec::new();
        content
            .try_push(
                "completion_message.content",
                ContentPart::Text(BoundedText::new("completion_message.text", text)?),
            )
            .map_err(LlmError::Boundary)?;
        Ok(Self {
            role,
            content,
            tool_call_id: None,
            proposed_calls: BoundedVec::new(),
        })
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum ToolSchema {
    Object {
        properties: BTreeMap<String, ToolSchema>,
        required: BTreeSet<String>,
        additional_properties: bool,
    },
    String,
    Integer,
    Number,
    Boolean,
    Array {
        items: Box<ToolSchema>,
        max_items: usize,
    },
    Nullable(Box<ToolSchema>),
}

impl ToolSchema {
    pub fn validate_schema(&self) -> Result<(), LlmError> {
        let mut remaining_nodes = MAX_SCHEMA_NODES;
        self.validate_schema_inner(0, &mut remaining_nodes)
    }

    fn validate_schema_inner(
        &self,
        depth: usize,
        remaining_nodes: &mut usize,
    ) -> Result<(), LlmError> {
        if depth > MAX_SCHEMA_DEPTH || *remaining_nodes == 0 {
            return Err(LlmError::InvalidToolSchema);
        }
        *remaining_nodes -= 1;
        match self {
            Self::Object {
                properties,
                required,
                ..
            } => {
                if properties.len() > MAX_SCHEMA_PROPERTIES
                    || required.len() > MAX_SCHEMA_PROPERTIES
                    || required.iter().any(|name| !properties.contains_key(name))
                    || properties
                        .keys()
                        .any(|name| name.is_empty() || name.len() > MAX_IDENTIFIER_BYTES)
                {
                    return Err(LlmError::InvalidToolSchema);
                }
                for schema in properties.values() {
                    schema.validate_schema_inner(depth + 1, remaining_nodes)?;
                }
                Ok(())
            }
            Self::Array { items, max_items } => {
                if *max_items == 0 || *max_items > 256 {
                    return Err(LlmError::InvalidToolSchema);
                }
                items.validate_schema_inner(depth + 1, remaining_nodes)
            }
            Self::Nullable(inner) => inner.validate_schema_inner(depth + 1, remaining_nodes),
            _ => Ok(()),
        }
    }

    pub fn validate_value(&self, value: &Value) -> Result<(), LlmError> {
        match self {
            Self::Object {
                properties,
                required,
                additional_properties,
            } => {
                let object = value.as_object().ok_or(LlmError::InvalidToolArguments)?;
                if required.iter().any(|name| !object.contains_key(name))
                    || (!additional_properties
                        && object.keys().any(|name| !properties.contains_key(name)))
                {
                    return Err(LlmError::InvalidToolArguments);
                }
                for (name, value) in object {
                    if let Some(schema) = properties.get(name) {
                        schema.validate_value(value)?;
                    }
                }
                Ok(())
            }
            Self::String if value.is_string() => Ok(()),
            Self::Integer if value.as_i64().is_some() || value.as_u64().is_some() => Ok(()),
            Self::Number if value.is_number() => Ok(()),
            Self::Boolean if value.is_boolean() => Ok(()),
            Self::Array { items, max_items } => {
                let values = value.as_array().ok_or(LlmError::InvalidToolArguments)?;
                if values.len() > *max_items {
                    return Err(LlmError::InvalidToolArguments);
                }
                for value in values {
                    items.validate_value(value)?;
                }
                Ok(())
            }
            Self::Nullable(_) if value.is_null() => Ok(()),
            Self::Nullable(inner) => inner.validate_value(value),
            _ => Err(LlmError::InvalidToolArguments),
        }
    }

    pub(crate) fn wire_value(&self) -> Value {
        match self {
            Self::Object {
                properties,
                required,
                additional_properties,
            } => {
                let properties = properties
                    .iter()
                    .map(|(name, schema)| (name.clone(), schema.wire_value()))
                    .collect::<Map<_, _>>();
                json!({
                    "type": "object",
                    "properties": properties,
                    "required": required,
                    "additionalProperties": additional_properties
                })
            }
            Self::String => json!({"type":"string"}),
            Self::Integer => json!({"type":"integer"}),
            Self::Number => json!({"type":"number"}),
            Self::Boolean => json!({"type":"boolean"}),
            Self::Array { items, max_items } => {
                json!({"type":"array", "items":items.wire_value(), "maxItems":max_items})
            }
            Self::Nullable(inner) => json!({"anyOf":[inner.wire_value(), {"type":"null"}]}),
        }
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ToolDefinition {
    pub name: BoundedText<MAX_IDENTIFIER_BYTES>,
    pub description: BoundedText<MAX_MESSAGE_BYTES>,
    pub schema: ToolSchema,
    pub mutating: bool,
}

impl ToolDefinition {
    pub fn validate(&self) -> Result<(), LlmError> {
        if !self
            .name
            .chars()
            .all(|character| character.is_ascii_alphanumeric() || matches!(character, '_' | '-'))
        {
            return Err(LlmError::InvalidToolSchema);
        }
        self.schema.validate_schema()
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum LlmError {
    Boundary(crate::types::BoundaryError),
    InvalidToolSchema,
    InvalidToolArguments,
}

impl fmt::Display for LlmError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Boundary(error) => write!(formatter, "LLM boundary rejected input: {error}"),
            Self::InvalidToolSchema => formatter.write_str("invalid registered tool schema"),
            Self::InvalidToolArguments => formatter.write_str("tool arguments do not match schema"),
        }
    }
}

impl Error for LlmError {}

impl From<crate::types::BoundaryError> for LlmError {
    fn from(value: crate::types::BoundaryError) -> Self {
        Self::Boundary(value)
    }
}

pub type MentraLlmProvider = mentra::provider_core::responses::ResponsesProvider<
    mentra::provider_core::StaticCredentialSource,
>;

pub struct LlmClient {
    connection: LlmConnection,
}

impl fmt::Debug for LlmClient {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("LlmClient")
            .field("endpoint", &self.connection.endpoint_url)
            .field("model", &self.connection.model)
            .finish_non_exhaustive()
    }
}

impl LlmClient {
    #[must_use]
    pub const fn new(connection: LlmConnection) -> Self {
        Self { connection }
    }

    #[must_use]
    pub fn mentra_provider(&self) -> MentraLlmProvider {
        let mut definition = mentra::provider_core::responses::openai_definition();
        definition.base_url = Some(self.connection.endpoint_url.expose_url().to_owned());
        mentra::provider_core::responses::ResponsesProvider::new(
            definition,
            mentra::provider_core::StaticCredentialSource::new(
                self.connection.api_key.expose_secret(),
            ),
        )
        .without_hybrid_http_previous_response_id()
    }

    #[must_use]
    pub fn configured_model(&self) -> &str {
        self.connection.model.as_deref().unwrap_or_default()
    }
}
