use async_trait::async_trait;
use crate::workflow::{WorkflowRepository, WorkflowStateMachine, MessageBroker};
use std::sync::Arc;
use tokio::sync::Mutex;
use dashmap::DashMap;
use uuid::Uuid;
use anyhow::Result;
use async_nats::Client;
use common_contracts::messages::WorkflowEvent;

// --- Real Implementation of Repository (In-Memory for MVP) ---
#[derive(Clone)]
pub struct DashMapWorkflowRepository {
    store: Arc<DashMap<Uuid, Arc<Mutex<WorkflowStateMachine>>>>
}

impl DashMapWorkflowRepository {
    pub fn new(store: Arc<DashMap<Uuid, Arc<Mutex<WorkflowStateMachine>>>>) -> Self {
        Self { store }
    }
}

#[async_trait]
impl WorkflowRepository for DashMapWorkflowRepository {
    async fn save(&self, workflow: &WorkflowStateMachine) -> Result<()> {
        if let Some(entry) = self.store.get(&workflow.request_id) {
            let mut guard = entry.lock().await;
            *guard = workflow.clone();
        } else {
            self.store.insert(workflow.request_id, Arc::new(Mutex::new(workflow.clone())));
        }
        Ok(())
    }

    async fn load(&self, id: Uuid) -> Result<Option<WorkflowStateMachine>> {
        if let Some(entry) = self.store.get(&id) {
            let guard = entry.lock().await;
            return Ok(Some(guard.clone()));
        }
        // Add debug log to see what keys are actually in the store
        // tracing::debug!("Keys in store: {:?}", self.store.iter().map(|r| *r.key()).collect::<Vec<_>>());
        Ok(None)
    }
}

// --- Real Implementation of Broker ---

#[derive(Clone)]
pub struct NatsMessageBroker {
    client: Client
}

impl NatsMessageBroker {
    pub fn new(client: Client) -> Self {
        Self { client }
    }
}

use common_contracts::subjects::NatsSubject;

#[async_trait]
impl MessageBroker for NatsMessageBroker {
    async fn publish_event(&self, request_id: Uuid, event: WorkflowEvent) -> Result<()> {
        let topic = NatsSubject::WorkflowProgress(request_id).to_string();
        let payload = serde_json::to_vec(&event)?;
        self.client.publish(topic, payload.into()).await.map_err(|e| anyhow::anyhow!(e))
    }

    async fn publish_command(&self, topic: &str, payload: Vec<u8>) -> Result<()> {
        self.client.publish(topic.to_string(), payload.into()).await.map_err(|e| anyhow::anyhow!(e))
    }
}
