use std::{
    collections::HashMap,
    sync::Arc,
    time::{Duration, Instant},
};

use crank_registry::{PostgresRegistry, PublishedAgentTool, RegistryError};
use tokio::sync::RwLock;
use tracing::info;

#[derive(Clone)]
pub struct PublishedToolCatalog {
    registry: PostgresRegistry,
    refresh_interval: Duration,
    cached: Arc<RwLock<HashMap<CatalogKey, CachedCatalog>>>,
}

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
struct CatalogKey {
    workspace_slug: String,
    agent_slug: String,
}

#[derive(Default)]
struct CachedCatalog {
    loaded_at: Option<Instant>,
    tools: Vec<PublishedAgentTool>,
}

impl PublishedToolCatalog {
    pub fn new(registry: PostgresRegistry, refresh_interval: Duration) -> Self {
        Self {
            registry,
            refresh_interval,
            cached: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    pub async fn list_tools(
        &self,
        workspace_slug: &str,
        agent_slug: &str,
    ) -> Result<Vec<PublishedAgentTool>, RegistryError> {
        self.refresh_if_stale(workspace_slug, agent_slug).await?;
        let guard = self.cached.read().await;
        Ok(guard
            .get(&CatalogKey::new(workspace_slug, agent_slug))
            .map(|entry| entry.tools.clone())
            .unwrap_or_default())
    }

    async fn refresh_if_stale(
        &self,
        workspace_slug: &str,
        agent_slug: &str,
    ) -> Result<(), RegistryError> {
        let key = CatalogKey::new(workspace_slug, agent_slug);
        let should_refresh = {
            let guard = self.cached.read().await;

            match guard.get(&key).and_then(|entry| entry.loaded_at) {
                Some(loaded_at) => loaded_at.elapsed() >= self.refresh_interval,
                None => true,
            }
        };

        if !should_refresh {
            return Ok(());
        }

        let tools = match self
            .registry
            .get_published_agent_tools_by_slug(workspace_slug, agent_slug)
            .await
        {
            Ok(tools) => tools,
            Err(RegistryError::PublishedAgentNotFound { .. }) => Vec::new(),
            Err(error) => return Err(error),
        };
        let mut guard = self.cached.write().await;
        let previous_count = guard
            .get(&key)
            .map(|entry| entry.tools.len())
            .unwrap_or_default();

        guard.insert(
            key,
            CachedCatalog {
                loaded_at: Some(Instant::now()),
                tools,
            },
        );

        info!(
            workspace_slug,
            agent_slug,
            published_tool_count = guard
                .get(&CatalogKey::new(workspace_slug, agent_slug))
                .map(|entry| entry.tools.len())
                .unwrap_or_default(),
            previous_published_tool_count = previous_count,
            "published agent catalog refreshed"
        );

        Ok(())
    }
}

impl CatalogKey {
    fn new(workspace_slug: &str, agent_slug: &str) -> Self {
        Self {
            workspace_slug: workspace_slug.to_owned(),
            agent_slug: agent_slug.to_owned(),
        }
    }
}