fastmcp-rs 0.2.0

use std::sync::Arc;

use async_trait::async_trait;
use chrono::{DateTime, Utc};
use indexmap::IndexMap;
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use tracing::{trace, warn};
use uuid::Uuid;

use crate::error::{FastMcpError, Result};
use crate::tool::DuplicateBehavior;

fn annotations_is_empty(map: &Map<String, Value>) -> bool {
    map.is_empty()
}

/// Wire representation of resource content.
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ResourceContent {
    Text {
        text: String,
        #[serde(default, skip_serializing_if = "Option::is_none")]
        mime_type: Option<String>,
    },
    Json {
        value: Value,
    },
}

impl ResourceContent {
    pub fn text(text: impl Into<String>) -> Self {
        Self::Text {
            text: text.into(),
            mime_type: Some("text/plain".into()),
        }
    }

    pub fn json(value: Value) -> Self {
        Self::Json { value }
    }
}

/// Metadata describing resources without loading content.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ResourceDefinitionMetadata {
    pub uri: String,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub name: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub description: Option<String>,
    #[serde(default, skip_serializing_if = "annotations_is_empty")]
    pub annotations: Map<String, Value>,
}

/// Runtime context for dynamic resource resolution.
#[derive(Clone, Debug)]
pub struct ResourceContext {
    pub uri: String,
    pub request_id: Uuid,
    pub timestamp: DateTime<Utc>,
}

impl ResourceContext {
    pub fn new(uri: impl Into<String>) -> Self {
        Self {
            uri: uri.into(),
            request_id: Uuid::new_v4(),
            timestamp: Utc::now(),
        }
    }
}

#[async_trait]
pub trait ResourceResolver: Send + Sync {
    async fn load(&self, ctx: ResourceContext) -> Result<ResourceContent>;
}

#[async_trait]
impl<F, Fut> ResourceResolver for F
where
    F: Send + Sync + Fn(ResourceContext) -> Fut,
    Fut: std::future::Future<Output = Result<ResourceContent>> + Send,
{
    async fn load(&self, ctx: ResourceContext) -> Result<ResourceContent> {
        (self)(ctx).await
    }
}

enum ResourceProvider {
    Static(ResourceContent),
    Dynamic(Arc<dyn ResourceResolver>),
}

pub struct ResourceDefinition {
    pub uri: String,
    pub name: Option<String>,
    pub description: Option<String>,
    pub annotations: Map<String, Value>,
    provider: ResourceProvider,
}

impl ResourceDefinition {
    pub fn static_resource(uri: impl Into<String>, content: ResourceContent) -> Self {
        Self {
            uri: uri.into(),
            name: None,
            description: None,
            annotations: Map::new(),
            provider: ResourceProvider::Static(content),
        }
    }

    pub fn dynamic(uri: impl Into<String>, resolver: impl ResourceResolver + 'static) -> Self {
        Self {
            uri: uri.into(),
            name: None,
            description: None,
            annotations: Map::new(),
            provider: ResourceProvider::Dynamic(Arc::new(resolver)),
        }
    }

    pub fn with_name(mut self, name: impl Into<String>) -> Self {
        self.name = Some(name.into());
        self
    }

    pub fn with_description(mut self, description: impl Into<String>) -> Self {
        self.description = Some(description.into());
        self
    }

    pub fn with_annotations(mut self, annotations: Map<String, Value>) -> Self {
        self.annotations = annotations;
        self
    }

    pub(crate) async fn load(&self) -> Result<ResourceContent> {
        match &self.provider {
            ResourceProvider::Static(content) => Ok(content.clone()),
            ResourceProvider::Dynamic(resolver) => {
                resolver.load(ResourceContext::new(self.uri.clone())).await
            }
        }
    }

    pub(crate) fn metadata(&self) -> ResourceDefinitionMetadata {
        ResourceDefinitionMetadata {
            uri: self.uri.clone(),
            name: self.name.clone(),
            description: self.description.clone(),
            annotations: self.annotations.clone(),
        }
    }
}

pub struct ResourceManager {
    duplicate_behavior: DuplicateBehavior,
    resources: RwLock<IndexMap<String, Arc<ResourceDefinition>>>,
}

impl ResourceManager {
    pub fn new(duplicate_behavior: DuplicateBehavior) -> Self {
        Self {
            duplicate_behavior,
            resources: RwLock::new(IndexMap::new()),
        }
    }

    pub fn register(&self, resource: ResourceDefinition) -> Result<()> {
        let mut guard = self.resources.write();
        match guard.get_mut(&resource.uri) {
            Some(existing) => match self.duplicate_behavior {
                DuplicateBehavior::Error => {
                    return Err(FastMcpError::DuplicateResource(resource.uri));
                }
                DuplicateBehavior::Ignore => {
                    trace!("Ignoring duplicate resource {}", resource.uri);
                }
                DuplicateBehavior::Replace => {
                    trace!("Replacing resource {}", resource.uri);
                    *existing = Arc::new(resource);
                }
                DuplicateBehavior::Warn => {
                    warn!("Replacing duplicate resource {}", resource.uri);
                    *existing = Arc::new(resource);
                }
            },
            None => {
                guard.insert(resource.uri.clone(), Arc::new(resource));
            }
        }
        Ok(())
    }

    pub fn list(&self) -> Vec<ResourceDefinitionMetadata> {
        self.resources
            .read()
            .values()
            .map(|res| res.metadata())
            .collect()
    }

    pub async fn read(&self, uri: &str) -> Result<ResourceContent> {
        let resource = {
            let guard = self.resources.read();
            guard
                .get(uri)
                .cloned()
                .ok_or_else(|| FastMcpError::ResourceNotFound(uri.to_string()))?
        };

        resource.load().await
    }
}

#[cfg(test)]
mod tests {
    use serde_json::json;

    use super::*;

    #[tokio::test]
    async fn static_resource_round_trip() {
        let manager = ResourceManager::new(DuplicateBehavior::Error);
        manager
            .register(
                ResourceDefinition::static_resource(
                    "resource://hello",
                    ResourceContent::json(json!({ "message": "ok" })),
                )
                .with_name("Hello resource"),
            )
            .unwrap();

        let listed = manager.list();
        assert_eq!(listed.len(), 1);
        assert_eq!(listed[0].uri, "resource://hello");

        let content = manager.read("resource://hello").await.unwrap();
        match content {
            ResourceContent::Json { value } => {
                assert_eq!(value["message"], "ok");
            }
            _ => panic!("expected json content"),
        }
    }
}