use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::debug;
use bincode;

/// Memory layer hierarchy
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum MemoryLayer {
    Global,
    Agent,
    Step,
    Tool,
}

/// Memory entry with metadata
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MemoryEntry {
    pub key: String,
    pub value: serde_json::Value,
    pub layer: MemoryLayer,
    pub timestamp: chrono::DateTime<chrono::Utc>,
    pub metadata: HashMap<String, String>,
}

/// Hierarchical memory system
pub struct Memory {
    /// Global memory (shared across all agents)
    global: Arc<RwLock<HashMap<String, MemoryEntry>>>,
    
    /// Agent-specific memory
    agents: Arc<RwLock<HashMap<String, HashMap<String, MemoryEntry>>>>,
    
    /// Step-specific memory (nested: agent_id -> step_id -> entries)
    steps: Arc<RwLock<HashMap<String, HashMap<String, HashMap<String, MemoryEntry>>>>>,
    
    /// Tool-specific memory (nested: agent_id -> step_id -> tool_id -> entries)
    tools: Arc<RwLock<HashMap<String, HashMap<String, HashMap<String, HashMap<String, MemoryEntry>>>>>>,
}

impl Memory {
    pub fn new() -> Self {
        Self {
            global: Arc::new(RwLock::new(HashMap::new())),
            agents: Arc::new(RwLock::new(HashMap::new())),
            steps: Arc::new(RwLock::new(HashMap::new())),
            tools: Arc::new(RwLock::new(HashMap::new())),
        }
    }
    
    /// Write to global memory
    pub async fn write_global(&self, key: String, value: serde_json::Value) {
        let entry = MemoryEntry {
            key: key.clone(),
            value,
            layer: MemoryLayer::Global,
            timestamp: chrono::Utc::now(),
            metadata: HashMap::new(),
        };
        
        let mut global = self.global.write().await;
        global.insert(key.clone(), entry);
        debug!("Wrote to global memory: {}", key);
    }
    
    /// Read from global memory
    pub async fn read_global(&self, key: &str) -> Option<serde_json::Value> {
        let global = self.global.read().await;
        global.get(key).map(|e| e.value.clone())
    }
    
    /// Write to agent memory
    pub async fn write_agent(&self, agent_id: &str, key: String, value: serde_json::Value) {
        let entry = MemoryEntry {
            key: key.clone(),
            value,
            layer: MemoryLayer::Agent,
            timestamp: chrono::Utc::now(),
            metadata: HashMap::new(),
        };
        
        let mut agents = self.agents.write().await;
        agents
            .entry(agent_id.to_string())
            .or_insert_with(HashMap::new)
            .insert(key.clone(), entry);
        debug!("Wrote to agent memory: {} -> {}", agent_id, key);
    }
    
    /// Read from agent memory
    pub async fn read_agent(&self, agent_id: &str, key: &str) -> Option<serde_json::Value> {
        let agents = self.agents.read().await;
        agents
            .get(agent_id)?
            .get(key)
            .map(|e| e.value.clone())
    }
    
    /// Write to step memory
    pub async fn write_step(
        &self,
        agent_id: &str,
        step_id: &str,
        key: String,
        value: serde_json::Value,
    ) {
        let entry = MemoryEntry {
            key: key.clone(),
            value,
            layer: MemoryLayer::Step,
            timestamp: chrono::Utc::now(),
            metadata: HashMap::new(),
        };
        
        let mut steps = self.steps.write().await;
        steps
            .entry(agent_id.to_string())
            .or_insert_with(HashMap::new)
            .entry(step_id.to_string())
            .or_insert_with(HashMap::new)
            .insert(key.clone(), entry);
        debug!("Wrote to step memory: {} -> {} -> {}", agent_id, step_id, key);
    }
    
    /// Read from step memory
    pub async fn read_step(
        &self,
        agent_id: &str,
        step_id: &str,
        key: &str,
    ) -> Option<serde_json::Value> {
        let steps = self.steps.read().await;
        steps
            .get(agent_id)?
            .get(step_id)?
            .get(key)
            .map(|e| e.value.clone())
    }
    
    /// Write to tool memory
    pub async fn write_tool(
        &self,
        agent_id: &str,
        step_id: &str,
        tool_id: &str,
        key: String,
        value: serde_json::Value,
    ) {
        let entry = MemoryEntry {
            key: key.clone(),
            value,
            layer: MemoryLayer::Tool,
            timestamp: chrono::Utc::now(),
            metadata: HashMap::new(),
        };
        
        let mut tools = self.tools.write().await;
        tools
            .entry(agent_id.to_string())
            .or_insert_with(HashMap::new)
            .entry(step_id.to_string())
            .or_insert_with(HashMap::new)
            .entry(tool_id.to_string())
            .or_insert_with(HashMap::new)
            .insert(key.clone(), entry);
        debug!(
            "Wrote to tool memory: {} -> {} -> {} -> {}",
            agent_id, step_id, tool_id, key
        );
    }
    
    /// Read from tool memory
    pub async fn read_tool(
        &self,
        agent_id: &str,
        step_id: &str,
        tool_id: &str,
        key: &str,
    ) -> Option<serde_json::Value> {
        let tools = self.tools.read().await;
        tools
            .get(agent_id)?
            .get(step_id)?
            .get(tool_id)?
            .get(key)
            .map(|e| e.value.clone())
    }
    
    /// Get all memory for an agent (hierarchical: global + agent + steps + tools)
    pub async fn get_agent_context(&self, agent_id: &str) -> HashMap<String, serde_json::Value> {
        let mut context = HashMap::new();
        
        // Add global memory
        {
            let global = self.global.read().await;
            for (key, entry) in global.iter() {
                context.insert(format!("global:{}", key), entry.value.clone());
            }
        }
        
        // Add agent memory
        {
            let agents = self.agents.read().await;
            if let Some(agent_mem) = agents.get(agent_id) {
                for (key, entry) in agent_mem.iter() {
                    context.insert(format!("agent:{}", key), entry.value.clone());
                }
            }
        }
        
        // Add step memory
        {
            let steps = self.steps.read().await;
            if let Some(agent_steps) = steps.get(agent_id) {
                for (step_id, step_mem) in agent_steps.iter() {
                    for (key, entry) in step_mem.iter() {
                        context.insert(
                            format!("step:{}:{}", step_id, key),
                            entry.value.clone(),
                        );
                    }
                }
            }
        }
        
        // Add tool memory
        {
            let tools = self.tools.read().await;
            if let Some(agent_tools) = tools.get(agent_id) {
                for (step_id, step_tools) in agent_tools.iter() {
                    for (tool_id, tool_mem) in step_tools.iter() {
                        for (key, entry) in tool_mem.iter() {
                            context.insert(
                                format!("tool:{}:{}:{}", step_id, tool_id, key),
                                entry.value.clone(),
                            );
                        }
                    }
                }
            }
        }
        
        context
    }
    
    /// Deterministic serialization of memory state
    pub async fn serialize(&self) -> Result<Vec<u8>, bincode::Error> {
        let snapshot = MemorySnapshot {
            global: {
                let global = self.global.read().await;
                global.clone()
            },
            agents: {
                let agents = self.agents.read().await;
                agents.clone()
            },
            steps: {
                let steps = self.steps.read().await;
                steps.clone()
            },
            tools: {
                let tools = self.tools.read().await;
                tools.clone()
            },
        };
        
        bincode::serialize(&snapshot)
    }
    
    /// Deterministic deserialization of memory state
    pub async fn deserialize(&self, data: &[u8]) -> Result<(), bincode::Error> {
        let snapshot: MemorySnapshot = bincode::deserialize(data)?;
        
        *self.global.write().await = snapshot.global;
        *self.agents.write().await = snapshot.agents;
        *self.steps.write().await = snapshot.steps;
        *self.tools.write().await = snapshot.tools;
        
        debug!("Deserialized memory snapshot");
        Ok(())
    }
    
    /// Clear all memory
    pub async fn clear(&self) {
        *self.global.write().await = HashMap::new();
        *self.agents.write().await = HashMap::new();
        *self.steps.write().await = HashMap::new();
        *self.tools.write().await = HashMap::new();
        debug!("Cleared all memory");
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct MemorySnapshot {
    global: HashMap<String, MemoryEntry>,
    agents: HashMap<String, HashMap<String, MemoryEntry>>,
    steps: HashMap<String, HashMap<String, HashMap<String, MemoryEntry>>>,
    tools: HashMap<String, HashMap<String, HashMap<String, HashMap<String, MemoryEntry>>>>,
}

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

