use std::{collections::HashMap, sync::Arc};

use tokio::sync::RwLock;
use uuid::Uuid;

#[derive(Clone)]
pub struct SessionStore {
    inner: Arc<RwLock<HashMap<String, SessionState>>>,
}

#[derive(Clone)]
pub struct SessionState {
    pub protocol_version: String,
    pub initialized: bool,
    pub workspace_slug: String,
    pub agent_slug: String,
}

impl SessionStore {
    pub fn new() -> Self {
        Self {
            inner: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    pub async fn create(
        &self,
        protocol_version: &str,
        workspace_slug: &str,
        agent_slug: &str,
    ) -> String {
        let session_id = Uuid::now_v7().to_string();
        let mut guard = self.inner.write().await;

        guard.insert(
            session_id.clone(),
            SessionState {
                protocol_version: protocol_version.to_owned(),
                initialized: false,
                workspace_slug: workspace_slug.to_owned(),
                agent_slug: agent_slug.to_owned(),
            },
        );

        session_id
    }

    pub async fn get(&self, session_id: &str) -> Option<SessionState> {
        let guard = self.inner.read().await;
        guard.get(session_id).cloned()
    }

    pub async fn mark_initialized(&self, session_id: &str) -> bool {
        let mut guard = self.inner.write().await;

        if let Some(session) = guard.get_mut(session_id) {
            session.initialized = true;
            return true;
        }

        false
    }

    pub async fn delete(&self, session_id: &str) -> bool {
        let mut guard = self.inner.write().await;
        guard.remove(session_id).is_some()
    }
}