pub mod anthropic;
pub mod coding;
pub mod communication;
pub mod instructions;
pub mod messages;
pub mod moose;
pub mod open_ai;
pub mod permissions;
pub mod tools;

use anthropic::AnthropicModel;
use anyhow::Result;
use communication::email::Email;
use messages::Message;
use open_ai::OpenAIModel;
use permissions::{PermissionStatus, ProgramPermissions};

use schemars::JsonSchema;
use serde::Serialize;
use tools::{Tool, ToolChoice, ToolDefinition, ToolExecution, Toolkit};

pub struct LLMContext {
    program_permissions: ProgramPermissions,
}

impl LLMContext {
    pub fn new(program_permissions: ProgramPermissions) -> Self {
        Self {
            program_permissions,
        }
    }

    pub fn check_program_permission(&self, program: &str) -> PermissionStatus {
        self.program_permissions.check_permission(program)
    }
}

/// LLM Provider Behaviors
///
/// Principle:
/// - Users can invoke the LLM with a set of messages
/// - LLM responds by adding a message to the conversation
/// - Users can stream the LLM's response
/// - Users can get a structured response from the LLM
pub trait LLMProvider<M: Message>: Send + Sync {
    /// Get a completion from the LLM
    fn get_completion(&self, messages: Vec<M>) -> Result<Vec<M>>;

    /// Stream a completion from the LLM
    fn stream_completion(&self, messages: Vec<M>) -> Result<Box<dyn Iterator<Item = M>>>;
}

/// Separate trait for structured responses
pub trait StructuredLLMProvider<M: Message>: LLMProvider<M> {
    /// Get a structured response from the LLM
    ///
    /// **DesiredSchema** is the type of the response to get from the LLM
    ///
    /// Example:
    /// ```
    /// #![derive(Serialize, Deserialize, JsonSchema)]
    /// struct MyStruct {
    ///     name: String,
    ///     age: u8,
    /// }
    /// ```
    fn get_structured_response<
        'de,
        DesiredSchema: Serialize + serde::de::DeserializeOwned + JsonSchema,
    >(
        &self,
        messages: Vec<M>,
    ) -> Result<DesiredSchema>;
}

/// LLM that can use a program on the user's computer
pub trait LLMComputerUser<M: Message> {
    fn can_use_program(&self, program: &str, llm_context: &LLMContext) -> PermissionStatus {
        // Default implementation checks if the program is allowed in the context's permissions
        llm_context.check_program_permission(program)
    }

    /// Use a program on the user's computer
    fn use_program(&self, program: &str, args: &str) -> Result<Vec<M>>;
}

pub trait LLMCommunication<M: Message> {
    fn send_email(&self, email: &Email) -> Result<Vec<M>>;
}

/// LLM that can use tools
///
/// How it all works:
/// - Ask the LLM to define a tool to be called. This is done wiht `get_chat_with_tools`
/// - When you have the tool that you want to use and it's arguments then you can run the tools
///
/// We provide functions to enable the user to easily have LLMs do work with tools.
///
/// If the user needs to customize the behavior when executing a tool, they can always get the
/// tool's execution object and execute it themselves by using the `get_chat_with_tools` function.
pub trait LLMToolUsage<M: Message> {
    /// Adds a next message to the conversation.
    ///
    /// If a tool was used the next messsage will include the tool's output
    fn get_work_result(
        &self,
        messages: Vec<M>,
        tool_kit: &Toolkit,
        tool_choice: &ToolChoice,
    ) -> Result<Vec<M>>;
    /// Adds a next message to the conversation with the result of a specific tool's execution
    fn do_work_with_tool(&self, messages: Vec<M>, tool: &dyn Tool) -> Result<Vec<M>>;
    /// Adds a next message to the conversation for subsequent tool use by the applicaiton runtime
    fn get_chat_with_tools(
        &self,
        messages: Vec<M>,
        tool_kit: &Toolkit,
        force_tool_use: &ToolChoice,
    ) -> Result<Vec<M>>;
}

/// LLM that can try to get a result, and retry if it fails
///
/// Principle:
/// - LLM tries to get a result
/// - If it fails, it can retry a given number of times
/// - If it fails again past the retry limit, it returns an error
///
pub trait LLMTrial<M: Message>: LLMProvider<M> {
    /// Attempt to get a result from the LLM
    ///
    /// *goal* is a prompt to the LLM to achieve the goal (ex: "Write a function to calculate the factorial of a number")
    /// *success* is a prompt to the LLM to check if the goal has been achieved (ex: "The function should take an integer and return its factorial")
    fn attempt(&self, goal: &str, success: &str, context: &LLMContext) -> Result<M>;

    /// Retry the LLM until it succeeds or the retry limit is reached
    fn retry(&self, attempts: u8, context: &LLMContext) -> Result<M>;
}

/// LLM Provider that support multiple models
///
/// Principle:
/// - Provider can be used with a model by using the `with_model` method
/// - Provider is created with a default model
/// - Provider enables getting the current model that's being used
pub trait MultiModelLLMProvider<Ms> {
    fn with_model(&self, model: Ms) -> Self;
    fn get_model(&self) -> Ms;
}

#[derive(Debug, Clone)]
pub enum ProviderModel {
    Anthropic(AnthropicModel),
    OpenAI(OpenAIModel),
}

impl ProviderModel {
    pub fn gpt_4o() -> Self {
        Self::OpenAI(OpenAIModel::Gpt4o)
    }

    pub fn o1_preview() -> Self {
        Self::OpenAI(OpenAIModel::O1Preview)
    }

    pub fn claude_3_5_sonnet() -> Self {
        Self::Anthropic(AnthropicModel::Claude35Sonnet)
    }
}
