pmat 3.16.0

#![cfg_attr(coverage_nightly, coverage(off))]

use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::HashMap;
use std::sync::Arc;

use super::types::McpError;

// ============================================================
// Tool Registry
// ============================================================

#[derive(Debug, Clone, Serialize, Deserialize)]
/// Metadata for tool.
pub struct ToolMetadata {
    pub name: String,
    pub description: String,
    pub input_schema: Value,
}

#[async_trait]
/// Trait defining Mcp tool behavior.
pub trait McpTool: Send + Sync {
    fn metadata(&self) -> ToolMetadata;
    async fn execute(&self, params: Value) -> Result<Value, McpError>;
}

/// Registry of tool instances.
pub struct ToolRegistry {
    pub(super) tools: HashMap<String, Arc<dyn McpTool>>,
    pub(super) metadata: HashMap<String, ToolMetadata>,
}

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

impl ToolRegistry {
    /// Create a new instance.
    pub fn new() -> Self {
        Self {
            tools: HashMap::new(),
            metadata: HashMap::new(),
        }
    }

    #[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
    /// Register a new item.
    pub fn register(&mut self, tool: Arc<dyn McpTool>) {
        let metadata = tool.metadata();
        self.tools.insert(metadata.name.clone(), tool);
        self.metadata.insert(metadata.name.clone(), metadata);
    }

    #[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
    /// List all registered tools.
    pub fn list(&self) -> Vec<ToolMetadata> {
        self.metadata.values().cloned().collect()
    }

    #[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
    /// Get a registered tool by name.
    pub fn get(&self, name: &str) -> Option<Arc<dyn McpTool>> {
        self.tools.get(name).cloned()
    }
}

// ============================================================
// Resource Registry
// ============================================================

#[derive(Debug, Clone, Serialize, Deserialize)]
/// Template for resource generation.
pub struct ResourceTemplate {
    pub uri_template: String,
    pub name: String,
    pub description: Option<String>,
    pub mime_type: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
/// Resource content.
pub struct ResourceContent {
    pub uri: String,
    pub mime_type: Option<String>,
    #[serde(flatten)]
    pub content: ResourceContentType,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
/// Type classification for resource content.
pub enum ResourceContentType {
    Text { text: String },
    Blob { blob: String }, // Base64 encoded
}

#[async_trait]
/// Trait defining Mcp resource behavior.
pub trait McpResource: Send + Sync {
    fn template(&self) -> ResourceTemplate;
    async fn read(&self, uri: &str) -> Result<ResourceContent, McpError>;
    fn subscribe(&self, uri: &str) -> Option<tokio::sync::watch::Receiver<ResourceContent>>;
}

/// Registry of resource instances.
pub struct ResourceRegistry {
    pub(super) resources: HashMap<String, Arc<dyn McpResource>>,
    pub(super) templates: HashMap<String, ResourceTemplate>,
}

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

impl ResourceRegistry {
    /// Create a new instance.
    pub fn new() -> Self {
        Self {
            resources: HashMap::new(),
            templates: HashMap::new(),
        }
    }

    #[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
    /// Register a new item.
    pub fn register(&mut self, resource: Arc<dyn McpResource>) {
        let template = resource.template();
        self.resources
            .insert(template.uri_template.clone(), resource);
        self.templates
            .insert(template.uri_template.clone(), template);
    }

    #[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
    /// List all registered resource templates.
    pub fn list(&self) -> Vec<ResourceTemplate> {
        self.templates.values().cloned().collect()
    }

    #[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
    /// Get a registered resource by URI template.
    pub fn get(&self, uri_template: &str) -> Option<Arc<dyn McpResource>> {
        self.resources.get(uri_template).cloned()
    }

    #[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
    /// Find matching.
    pub fn find_matching(&self, uri: &str) -> Option<Arc<dyn McpResource>> {
        // Simple pattern matching - could be enhanced
        for (template, resource) in &self.resources {
            if uri.starts_with(&template.replace("{}", "")) {
                return Some(resource.clone());
            }
        }
        None
    }
}

// ============================================================
// Prompt Registry
// ============================================================

#[derive(Debug, Clone, Serialize, Deserialize)]
/// Metadata for prompt.
pub struct PromptMetadata {
    pub name: String,
    pub description: Option<String>,
    pub arguments: Option<Vec<PromptArgument>>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
/// Prompt argument.
pub struct PromptArgument {
    pub name: String,
    pub description: Option<String>,
    pub required: Option<bool>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
/// Prompt message.
pub struct PromptMessage {
    pub role: String,
    pub content: PromptContent,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
/// Content variants for prompt.
pub enum PromptContent {
    Text(String),
    Parts(Vec<ContentPart>),
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type")]
/// Content part variants for content.
pub enum ContentPart {
    #[serde(rename = "text")]
    Text { text: String },
    #[serde(rename = "image")]
    Image { data: String, mime_type: String },
    #[serde(rename = "resource")]
    Resource { uri: String },
}

#[async_trait]
/// Trait defining Mcp prompt behavior.
pub trait McpPrompt: Send + Sync {
    fn metadata(&self) -> PromptMetadata;
    async fn get(
        &self,
        arguments: Option<HashMap<String, String>>,
    ) -> Result<Vec<PromptMessage>, McpError>;
}

/// Registry of prompt instances.
pub struct PromptRegistry {
    pub(super) prompts: HashMap<String, Arc<dyn McpPrompt>>,
    pub(super) metadata: HashMap<String, PromptMetadata>,
}

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

impl PromptRegistry {
    /// Create a new instance.
    pub fn new() -> Self {
        Self {
            prompts: HashMap::new(),
            metadata: HashMap::new(),
        }
    }

    #[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
    /// Register a new item.
    pub fn register(&mut self, prompt: Arc<dyn McpPrompt>) {
        let metadata = prompt.metadata();
        self.prompts.insert(metadata.name.clone(), prompt);
        self.metadata.insert(metadata.name.clone(), metadata);
    }

    #[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
    /// List all registered prompts.
    pub fn list(&self) -> Vec<PromptMetadata> {
        self.metadata.values().cloned().collect()
    }

    #[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
    /// Get a registered prompt by name.
    pub fn get(&self, name: &str) -> Option<Arc<dyn McpPrompt>> {
        self.prompts.get(name).cloned()
    }
}