//! Graph storage engine.

use super::edge::{Edge, EdgeId};
use super::node::{Node, NodeId};
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use serde_json::Value as JsonValue;
use std::collections::{HashMap, HashSet};
use thiserror::Error;

/// Direction for traversing edges.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Direction {
    /// Follow outgoing edges only.
    Outgoing,
    /// Follow incoming edges only.
    Incoming,
    /// Follow edges in both directions.
    Both,
}

/// Graph storage errors.
#[derive(Debug, Error)]
pub enum GraphError {
    /// Node not found.
    #[error("Node not found: {0}")]
    NodeNotFound(String),

    /// Edge not found.
    #[error("Edge not found: {0}")]
    EdgeNotFound(String),

    /// Node already exists.
    #[error("Node already exists: {0}")]
    NodeExists(String),

    /// Invalid operation.
    #[error("Invalid operation: {0}")]
    InvalidOperation(String),

    /// Constraint violation.
    #[error("Constraint violation: {0}")]
    ConstraintViolation(String),
}

/// Statistics for a graph store.
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
pub struct GraphStats {
    /// Total number of nodes.
    pub node_count: u64,
    /// Total number of edges.
    pub edge_count: u64,
    /// Number of distinct labels.
    pub label_count: u64,
    /// Number of distinct edge types.
    pub edge_type_count: u64,
}

/// Graph storage engine.
pub struct GraphStore {
    /// Node storage.
    nodes: RwLock<HashMap<NodeId, Node>>,
    /// Edge storage.
    edges: RwLock<HashMap<EdgeId, Edge>>,
    /// Outgoing adjacency list: node -> outgoing edges.
    adjacency: RwLock<HashMap<NodeId, Vec<EdgeId>>>,
    /// Incoming adjacency list: node -> incoming edges.
    reverse_adj: RwLock<HashMap<NodeId, Vec<EdgeId>>>,
    /// Label index: label -> nodes with that label.
    label_index: RwLock<HashMap<String, HashSet<NodeId>>>,
    /// Edge type index: type -> edges of that type.
    edge_type_index: RwLock<HashMap<String, HashSet<EdgeId>>>,
}

impl GraphStore {
    /// Creates a new empty graph store.
    pub fn new() -> Self {
        Self {
            nodes: RwLock::new(HashMap::new()),
            edges: RwLock::new(HashMap::new()),
            adjacency: RwLock::new(HashMap::new()),
            reverse_adj: RwLock::new(HashMap::new()),
            label_index: RwLock::new(HashMap::new()),
            edge_type_index: RwLock::new(HashMap::new()),
        }
    }

    /// Returns statistics about the graph.
    pub fn stats(&self) -> GraphStats {
        GraphStats {
            node_count: self.nodes.read().len() as u64,
            edge_count: self.edges.read().len() as u64,
            label_count: self.label_index.read().len() as u64,
            edge_type_count: self.edge_type_index.read().len() as u64,
        }
    }

    // ==================== Node Operations ====================

    /// Creates a new node with the given labels and properties.
    pub fn create_node(&self, labels: Vec<String>, properties: JsonValue) -> NodeId {
        let node = Node::new(labels.clone(), properties);
        let id = node.id;

        // Update label index
        {
            let mut label_idx = self.label_index.write();
            for label in &labels {
                label_idx.entry(label.clone()).or_default().insert(id);
            }
        }

        // Initialize adjacency lists
        self.adjacency.write().insert(id, Vec::new());
        self.reverse_adj.write().insert(id, Vec::new());

        // Store node
        self.nodes.write().insert(id, node);

        id
    }

    /// Gets a node by ID.
    pub fn get_node(&self, id: &NodeId) -> Option<Node> {
        self.nodes.read().get(id).cloned()
    }

    /// Updates a node's properties.
    pub fn update_node(&self, id: &NodeId, properties: JsonValue) -> Result<(), GraphError> {
        let mut nodes = self.nodes.write();
        let node = nodes
            .get_mut(id)
            .ok_or_else(|| GraphError::NodeNotFound(id.to_string()))?;

        node.properties = properties;
        node.updated_at = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_millis() as u64;

        Ok(())
    }

    /// Updates a node's labels.
    pub fn update_node_labels(&self, id: &NodeId, labels: Vec<String>) -> Result<(), GraphError> {
        let mut nodes = self.nodes.write();
        let node = nodes
            .get_mut(id)
            .ok_or_else(|| GraphError::NodeNotFound(id.to_string()))?;

        // Update label index
        {
            let mut label_idx = self.label_index.write();

            // Remove old labels
            for old_label in &node.labels {
                if let Some(set) = label_idx.get_mut(old_label) {
                    set.remove(id);
                }
            }

            // Add new labels
            for new_label in &labels {
                label_idx.entry(new_label.clone()).or_default().insert(*id);
            }
        }

        node.labels = labels;
        node.updated_at = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_millis() as u64;

        Ok(())
    }

    /// Deletes a node and all its connected edges.
    pub fn delete_node(&self, id: &NodeId) -> Result<(), GraphError> {
        // Get connected edges
        let outgoing: Vec<EdgeId> = self
            .adjacency
            .read()
            .get(id)
            .cloned()
            .unwrap_or_default();
        let incoming: Vec<EdgeId> = self
            .reverse_adj
            .read()
            .get(id)
            .cloned()
            .unwrap_or_default();

        // Delete all connected edges
        for edge_id in outgoing.iter().chain(incoming.iter()) {
            let _ = self.delete_edge(edge_id);
        }

        // Remove from label index
        {
            let nodes = self.nodes.read();
            if let Some(node) = nodes.get(id) {
                let mut label_idx = self.label_index.write();
                for label in &node.labels {
                    if let Some(set) = label_idx.get_mut(label) {
                        set.remove(id);
                    }
                }
            }
        }

        // Remove adjacency entries
        self.adjacency.write().remove(id);
        self.reverse_adj.write().remove(id);

        // Remove node
        self.nodes
            .write()
            .remove(id)
            .ok_or_else(|| GraphError::NodeNotFound(id.to_string()))?;

        Ok(())
    }

    /// Finds nodes by label.
    pub fn find_nodes_by_label(&self, label: &str) -> Vec<Node> {
        let node_ids: Vec<NodeId> = self
            .label_index
            .read()
            .get(label)
            .map(|set| set.iter().cloned().collect())
            .unwrap_or_default();

        let nodes = self.nodes.read();
        node_ids
            .iter()
            .filter_map(|id| nodes.get(id).cloned())
            .collect()
    }

    /// Finds nodes matching a filter.
    pub fn find_nodes(&self, label: Option<&str>, filter: &JsonValue) -> Vec<Node> {
        let candidates: Vec<Node> = if let Some(l) = label {
            self.find_nodes_by_label(l)
        } else {
            self.nodes.read().values().cloned().collect()
        };

        candidates
            .into_iter()
            .filter(|n| n.matches_properties(filter))
            .collect()
    }

    // ==================== Edge Operations ====================

    /// Creates a new edge.
    pub fn create_edge(
        &self,
        source: NodeId,
        target: NodeId,
        edge_type: impl Into<String>,
        properties: JsonValue,
    ) -> Result<EdgeId, GraphError> {
        let edge_type = edge_type.into();

        // Verify both nodes exist
        {
            let nodes = self.nodes.read();
            if !nodes.contains_key(&source) {
                return Err(GraphError::NodeNotFound(source.to_string()));
            }
            if !nodes.contains_key(&target) {
                return Err(GraphError::NodeNotFound(target.to_string()));
            }
        }

        let edge = Edge::new(source, target, edge_type.clone(), properties);
        let id = edge.id;

        // Update adjacency lists
        self.adjacency.write().entry(source).or_default().push(id);
        self.reverse_adj.write().entry(target).or_default().push(id);

        // Update edge type index
        self.edge_type_index
            .write()
            .entry(edge_type)
            .or_default()
            .insert(id);

        // Store edge
        self.edges.write().insert(id, edge);

        Ok(id)
    }

    /// Creates an undirected edge.
    pub fn create_undirected_edge(
        &self,
        source: NodeId,
        target: NodeId,
        edge_type: impl Into<String>,
        properties: JsonValue,
    ) -> Result<EdgeId, GraphError> {
        let edge_type = edge_type.into();

        // Verify both nodes exist
        {
            let nodes = self.nodes.read();
            if !nodes.contains_key(&source) {
                return Err(GraphError::NodeNotFound(source.to_string()));
            }
            if !nodes.contains_key(&target) {
                return Err(GraphError::NodeNotFound(target.to_string()));
            }
        }

        let edge = Edge::undirected(source, target, edge_type.clone(), properties);
        let id = edge.id;

        // Update adjacency lists (both directions for undirected)
        {
            let mut adj = self.adjacency.write();
            adj.entry(source).or_default().push(id);
            adj.entry(target).or_default().push(id);
        }
        {
            let mut rev = self.reverse_adj.write();
            rev.entry(source).or_default().push(id);
            rev.entry(target).or_default().push(id);
        }

        // Update edge type index
        self.edge_type_index
            .write()
            .entry(edge_type)
            .or_default()
            .insert(id);

        // Store edge
        self.edges.write().insert(id, edge);

        Ok(id)
    }

    /// Gets an edge by ID.
    pub fn get_edge(&self, id: &EdgeId) -> Option<Edge> {
        self.edges.read().get(id).cloned()
    }

    /// Deletes an edge.
    pub fn delete_edge(&self, id: &EdgeId) -> Result<(), GraphError> {
        let edge = self
            .edges
            .write()
            .remove(id)
            .ok_or_else(|| GraphError::EdgeNotFound(id.to_string()))?;

        // Update adjacency lists
        {
            let mut adj = self.adjacency.write();
            if let Some(list) = adj.get_mut(&edge.source) {
                list.retain(|e| e != id);
            }
            if !edge.directed {
                if let Some(list) = adj.get_mut(&edge.target) {
                    list.retain(|e| e != id);
                }
            }
        }
        {
            let mut rev = self.reverse_adj.write();
            if let Some(list) = rev.get_mut(&edge.target) {
                list.retain(|e| e != id);
            }
            if !edge.directed {
                if let Some(list) = rev.get_mut(&edge.source) {
                    list.retain(|e| e != id);
                }
            }
        }

        // Update edge type index
        if let Some(set) = self.edge_type_index.write().get_mut(&edge.edge_type) {
            set.remove(id);
        }

        Ok(())
    }

    /// Finds edges by type.
    pub fn find_edges_by_type(&self, edge_type: &str) -> Vec<Edge> {
        let edge_ids: Vec<EdgeId> = self
            .edge_type_index
            .read()
            .get(edge_type)
            .map(|set| set.iter().cloned().collect())
            .unwrap_or_default();

        let edges = self.edges.read();
        edge_ids
            .iter()
            .filter_map(|id| edges.get(id).cloned())
            .collect()
    }

    // ==================== Traversal Operations ====================

    /// Gets neighboring nodes.
    pub fn neighbors(&self, id: &NodeId, direction: Direction) -> Vec<Node> {
        let edge_ids = self.edges_of_node(id, direction);
        let edges = self.edges.read();
        let nodes = self.nodes.read();

        let mut neighbor_ids = HashSet::new();
        for eid in edge_ids {
            if let Some(edge) = edges.get(&eid) {
                if let Some(other) = self.get_neighbor_from_edge(edge, id, direction) {
                    neighbor_ids.insert(other);
                }
            }
        }

        neighbor_ids
            .iter()
            .filter_map(|nid| nodes.get(nid).cloned())
            .collect()
    }

    /// Gets edges connected to a node.
    pub fn edges_of(&self, id: &NodeId, direction: Direction) -> Vec<Edge> {
        let edge_ids = self.edges_of_node(id, direction);
        let edges = self.edges.read();

        edge_ids
            .iter()
            .filter_map(|eid| edges.get(eid).cloned())
            .collect()
    }

    /// Gets edge IDs connected to a node.
    fn edges_of_node(&self, id: &NodeId, direction: Direction) -> Vec<EdgeId> {
        match direction {
            Direction::Outgoing => self.adjacency.read().get(id).cloned().unwrap_or_default(),
            Direction::Incoming => self.reverse_adj.read().get(id).cloned().unwrap_or_default(),
            Direction::Both => {
                let mut result = self.adjacency.read().get(id).cloned().unwrap_or_default();
                let incoming = self.reverse_adj.read().get(id).cloned().unwrap_or_default();
                for eid in incoming {
                    if !result.contains(&eid) {
                        result.push(eid);
                    }
                }
                result
            }
        }
    }

    /// Gets the neighbor node from an edge.
    fn get_neighbor_from_edge(&self, edge: &Edge, from: &NodeId, direction: Direction) -> Option<NodeId> {
        match direction {
            Direction::Outgoing => {
                if &edge.source == from {
                    Some(edge.target)
                } else if !edge.directed && &edge.target == from {
                    Some(edge.source)
                } else {
                    None
                }
            }
            Direction::Incoming => {
                if &edge.target == from {
                    Some(edge.source)
                } else if !edge.directed && &edge.source == from {
                    Some(edge.target)
                } else {
                    None
                }
            }
            Direction::Both => edge.other_end(from).or_else(|| {
                // For directed edges, still return the other end
                if &edge.source == from {
                    Some(edge.target)
                } else if &edge.target == from {
                    Some(edge.source)
                } else {
                    None
                }
            }),
        }
    }

    /// Gets neighbors connected by a specific edge type.
    pub fn neighbors_by_type(&self, id: &NodeId, edge_type: &str, direction: Direction) -> Vec<Node> {
        let edges = self.edges_of(id, direction);
        let nodes = self.nodes.read();

        let mut neighbor_ids = HashSet::new();
        for edge in edges {
            if edge.edge_type == edge_type {
                if let Some(other) = self.get_neighbor_from_edge(&edge, id, direction) {
                    neighbor_ids.insert(other);
                }
            }
        }

        neighbor_ids
            .iter()
            .filter_map(|nid| nodes.get(nid).cloned())
            .collect()
    }

    /// Checks if an edge exists between two nodes.
    pub fn has_edge(&self, source: &NodeId, target: &NodeId, edge_type: Option<&str>) -> bool {
        let edges = self.edges_of(source, Direction::Outgoing);
        for edge in edges {
            if &edge.target == target {
                if let Some(et) = edge_type {
                    if edge.edge_type == et {
                        return true;
                    }
                } else {
                    return true;
                }
            }
        }
        false
    }

    /// Gets all edges between two nodes.
    pub fn edges_between(&self, source: &NodeId, target: &NodeId) -> Vec<Edge> {
        let edges = self.edges_of(source, Direction::Outgoing);
        edges
            .into_iter()
            .filter(|e| &e.target == target || (!e.directed && &e.source == target))
            .collect()
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_create_node() {
        let store = GraphStore::new();

        let id = store.create_node(
            vec!["User".to_string()],
            serde_json::json!({"name": "Alice"}),
        );

        let node = store.get_node(&id).unwrap();
        assert!(node.has_label("User"));
        assert_eq!(node.get_property("name"), Some(&serde_json::json!("Alice")));
    }

    #[test]
    fn test_create_edge() {
        let store = GraphStore::new();

        let alice = store.create_node(vec!["User".to_string()], serde_json::json!({"name": "Alice"}));
        let bob = store.create_node(vec!["User".to_string()], serde_json::json!({"name": "Bob"}));

        let edge_id = store
            .create_edge(alice, bob, "FRIEND", serde_json::json!({"since": 2020}))
            .unwrap();

        let edge = store.get_edge(&edge_id).unwrap();
        assert_eq!(edge.source, alice);
        assert_eq!(edge.target, bob);
        assert_eq!(edge.edge_type, "FRIEND");
    }

    #[test]
    fn test_neighbors() {
        let store = GraphStore::new();

        let alice = store.create_node(vec!["User".to_string()], serde_json::json!({"name": "Alice"}));
        let bob = store.create_node(vec!["User".to_string()], serde_json::json!({"name": "Bob"}));
        let charlie = store.create_node(vec!["User".to_string()], serde_json::json!({"name": "Charlie"}));

        store.create_edge(alice, bob, "FRIEND", serde_json::json!({})).unwrap();
        store.create_edge(alice, charlie, "FRIEND", serde_json::json!({})).unwrap();

        let neighbors = store.neighbors(&alice, Direction::Outgoing);
        assert_eq!(neighbors.len(), 2);
    }

    #[test]
    fn test_find_by_label() {
        let store = GraphStore::new();

        store.create_node(vec!["User".to_string()], serde_json::json!({"name": "Alice"}));
        store.create_node(vec!["User".to_string()], serde_json::json!({"name": "Bob"}));
        store.create_node(vec!["Product".to_string()], serde_json::json!({"name": "Widget"}));

        let users = store.find_nodes_by_label("User");
        assert_eq!(users.len(), 2);

        let products = store.find_nodes_by_label("Product");
        assert_eq!(products.len(), 1);
    }

    #[test]
    fn test_delete_node() {
        let store = GraphStore::new();

        let alice = store.create_node(vec!["User".to_string()], serde_json::json!({}));
        let bob = store.create_node(vec!["User".to_string()], serde_json::json!({}));

        store.create_edge(alice, bob, "FRIEND", serde_json::json!({})).unwrap();

        // Delete Alice - should also delete the edge
        store.delete_node(&alice).unwrap();

        assert!(store.get_node(&alice).is_none());
        assert_eq!(store.stats().edge_count, 0);
    }

    #[test]
    fn test_undirected_edge() {
        let store = GraphStore::new();

        let a = store.create_node(vec![], serde_json::json!({}));
        let b = store.create_node(vec![], serde_json::json!({}));

        store.create_undirected_edge(a, b, "LINK", serde_json::json!({})).unwrap();

        // Both directions should work
        let a_neighbors = store.neighbors(&a, Direction::Outgoing);
        let b_neighbors = store.neighbors(&b, Direction::Outgoing);

        assert_eq!(a_neighbors.len(), 1);
        assert_eq!(b_neighbors.len(), 1);
    }

    #[test]
    fn test_edges_between() {
        let store = GraphStore::new();

        let a = store.create_node(vec![], serde_json::json!({}));
        let b = store.create_node(vec![], serde_json::json!({}));

        store.create_edge(a, b, "TYPE_A", serde_json::json!({})).unwrap();
        store.create_edge(a, b, "TYPE_B", serde_json::json!({})).unwrap();

        let edges = store.edges_between(&a, &b);
        assert_eq!(edges.len(), 2);
    }
}
