//! Graph storage for knowledge graph entities and relations
//!
//! Provides persistent storage and querying for the knowledge graph
//! with efficient traversal operations.

use std::fs::{self, File};
use std::io::{self, BufRead, BufReader, BufWriter, Write};
use std::path::{Path, PathBuf};
use std::collections::{HashMap, HashSet};

use parking_lot::RwLock;

use crate::types::{Entity, Relation};

/// Graph storage for entities and relations
pub struct GraphStore {
    /// Path to entities file
    entities_path: PathBuf,
    
    /// Path to relations file
    relations_path: PathBuf,
    
    /// In-memory entity storage
    entities: RwLock<HashMap<String, Entity>>,
    
    /// In-memory relation storage
    relations: RwLock<HashMap<String, Relation>>,
    
    /// Index: entity_id -> outgoing relation_ids
    outgoing_relations: RwLock<HashMap<String, Vec<String>>>,
    
    /// Index: entity_id -> incoming relation_ids
    incoming_relations: RwLock<HashMap<String, Vec<String>>>,
    
    /// Index: entity_type -> entity_ids
    entity_type_index: RwLock<HashMap<String, Vec<String>>>,
    
    /// Whether store has unsaved changes
    dirty: RwLock<bool>,
}

impl GraphStore {
    /// Create a new graph store at the given directory
    pub fn new(dir: impl AsRef<Path>) -> io::Result<Self> {
        let dir = dir.as_ref();
        fs::create_dir_all(dir)?;
        
        let store = Self {
            entities_path: dir.join("entities.jsonl"),
            relations_path: dir.join("relations.jsonl"),
            entities: RwLock::new(HashMap::new()),
            relations: RwLock::new(HashMap::new()),
            outgoing_relations: RwLock::new(HashMap::new()),
            incoming_relations: RwLock::new(HashMap::new()),
            entity_type_index: RwLock::new(HashMap::new()),
            dirty: RwLock::new(false),
        };
        
        Ok(store)
    }
    
    /// Open an existing store or create a new one
    pub fn open_or_create(dir: impl AsRef<Path>) -> io::Result<Self> {
        let mut store = Self::new(dir)?;
        
        if store.entities_path.exists() {
            store.load()?;
        }
        
        Ok(store)
    }
    
    /// Load graph data from files
    fn load(&mut self) -> io::Result<()> {
        // Load entities
        if self.entities_path.exists() {
            let file = File::open(&self.entities_path)?;
            let reader = BufReader::new(file);
            
            let mut entities = HashMap::new();
            let mut type_index: HashMap<String, Vec<String>> = HashMap::new();
            
            for line in reader.lines() {
                let line = line?;
                if line.trim().is_empty() {
                    continue;
                }
                
                let entity: Entity = serde_json::from_str(&line)
                    .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
                
                type_index.entry(entity.entity_type.clone())
                    .or_default()
                    .push(entity.id.clone());
                
                entities.insert(entity.id.clone(), entity);
            }
            
            *self.entities.write() = entities;
            *self.entity_type_index.write() = type_index;
        }
        
        // Load relations
        if self.relations_path.exists() {
            let file = File::open(&self.relations_path)?;
            let reader = BufReader::new(file);
            
            let mut relations = HashMap::new();
            let mut outgoing: HashMap<String, Vec<String>> = HashMap::new();
            let mut incoming: HashMap<String, Vec<String>> = HashMap::new();
            
            for line in reader.lines() {
                let line = line?;
                if line.trim().is_empty() {
                    continue;
                }
                
                let relation: Relation = serde_json::from_str(&line)
                    .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
                
                outgoing.entry(relation.source.clone())
                    .or_default()
                    .push(relation.id.clone());
                
                incoming.entry(relation.target.clone())
                    .or_default()
                    .push(relation.id.clone());
                
                relations.insert(relation.id.clone(), relation);
            }
            
            *self.relations.write() = relations;
            *self.outgoing_relations.write() = outgoing;
            *self.incoming_relations.write() = incoming;
        }
        
        *self.dirty.write() = false;
        Ok(())
    }
    
    /// Save graph data to files
    pub fn save(&self) -> io::Result<()> {
        // Save entities
        {
            let file = File::create(&self.entities_path)?;
            let mut writer = BufWriter::new(file);
            
            for entity in self.entities.read().values() {
                let json = serde_json::to_string(entity)
                    .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
                writeln!(writer, "{}", json)?;
            }
            
            writer.flush()?;
        }
        
        // Save relations
        {
            let file = File::create(&self.relations_path)?;
            let mut writer = BufWriter::new(file);
            
            for relation in self.relations.read().values() {
                let json = serde_json::to_string(relation)
                    .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
                writeln!(writer, "{}", json)?;
            }
            
            writer.flush()?;
        }
        
        *self.dirty.write() = false;
        Ok(())
    }
    
    // === Entity Operations ===
    
    /// Add or update an entity
    pub fn put_entity(&self, entity: Entity) {
        let mut entities = self.entities.write();
        let mut type_index = self.entity_type_index.write();
        
        // Remove from old type index if updating
        if let Some(old) = entities.get(&entity.id) {
            if old.entity_type != entity.entity_type {
                if let Some(ids) = type_index.get_mut(&old.entity_type) {
                    ids.retain(|id| id != &entity.id);
                }
            }
        }
        
        // Add to type index
        type_index.entry(entity.entity_type.clone())
            .or_default()
            .push(entity.id.clone());
        
        entities.insert(entity.id.clone(), entity);
        *self.dirty.write() = true;
    }
    
    /// Get an entity by ID
    pub fn get_entity(&self, id: &str) -> Option<Entity> {
        self.entities.read().get(id).cloned()
    }
    
    /// Delete an entity by ID
    pub fn delete_entity(&self, id: &str) -> Option<Entity> {
        let mut entities = self.entities.write();
        let mut type_index = self.entity_type_index.write();
        
        let removed = entities.remove(id);
        
        if let Some(ref entity) = removed {
            // Remove from type index
            if let Some(ids) = type_index.get_mut(&entity.entity_type) {
                ids.retain(|eid| eid != id);
            }
            
            // Note: Relations referencing this entity should also be cleaned up
            // by the caller
            
            *self.dirty.write() = true;
        }
        
        removed
    }
    
    /// Check if entity exists
    pub fn contains_entity(&self, id: &str) -> bool {
        self.entities.read().contains_key(id)
    }
    
    /// Get all entities of a given type
    pub fn find_entities_by_type(&self, entity_type: &str) -> Vec<Entity> {
        let type_index = self.entity_type_index.read();
        let entities = self.entities.read();
        
        type_index.get(entity_type)
            .map(|ids| {
                ids.iter()
                    .filter_map(|id| entities.get(id).cloned())
                    .collect()
            })
            .unwrap_or_default()
    }
    
    /// Get all entities
    pub fn all_entities(&self) -> Vec<Entity> {
        self.entities.read().values().cloned().collect()
    }
    
    /// Get entity count
    pub fn entity_count(&self) -> usize {
        self.entities.read().len()
    }
    
    // === Relation Operations ===
    
    /// Add or update a relation
    pub fn put_relation(&self, relation: Relation) {
        let mut relations = self.relations.write();
        let mut outgoing = self.outgoing_relations.write();
        let mut incoming = self.incoming_relations.write();
        
        // Remove old indices if updating
        if let Some(old) = relations.get(&relation.id) {
            if let Some(ids) = outgoing.get_mut(&old.source) {
                ids.retain(|id| id != &relation.id);
            }
            if let Some(ids) = incoming.get_mut(&old.target) {
                ids.retain(|id| id != &relation.id);
            }
        }
        
        // Add to indices
        outgoing.entry(relation.source.clone())
            .or_default()
            .push(relation.id.clone());
        
        incoming.entry(relation.target.clone())
            .or_default()
            .push(relation.id.clone());
        
        relations.insert(relation.id.clone(), relation);
        *self.dirty.write() = true;
    }
    
    /// Get a relation by ID
    pub fn get_relation(&self, id: &str) -> Option<Relation> {
        self.relations.read().get(id).cloned()
    }
    
    /// Delete a relation by ID
    pub fn delete_relation(&self, id: &str) -> Option<Relation> {
        let mut relations = self.relations.write();
        let mut outgoing = self.outgoing_relations.write();
        let mut incoming = self.incoming_relations.write();
        
        let removed = relations.remove(id);
        
        if let Some(ref relation) = removed {
            if let Some(ids) = outgoing.get_mut(&relation.source) {
                ids.retain(|rid| rid != id);
            }
            if let Some(ids) = incoming.get_mut(&relation.target) {
                ids.retain(|rid| rid != id);
            }
            
            *self.dirty.write() = true;
        }
        
        removed
    }
    
    /// Get all relations
    pub fn all_relations(&self) -> Vec<Relation> {
        self.relations.read().values().cloned().collect()
    }
    
    /// Get relation count
    pub fn relation_count(&self) -> usize {
        self.relations.read().len()
    }
    
    // === Graph Traversal ===
    
    /// Get outgoing relations from an entity
    pub fn outgoing_from(&self, entity_id: &str) -> Vec<Relation> {
        let outgoing = self.outgoing_relations.read();
        let relations = self.relations.read();
        
        outgoing.get(entity_id)
            .map(|ids| {
                ids.iter()
                    .filter_map(|id| relations.get(id).cloned())
                    .collect()
            })
            .unwrap_or_default()
    }
    
    /// Get incoming relations to an entity
    pub fn incoming_to(&self, entity_id: &str) -> Vec<Relation> {
        let incoming = self.incoming_relations.read();
        let relations = self.relations.read();
        
        incoming.get(entity_id)
            .map(|ids| {
                ids.iter()
                    .filter_map(|id| relations.get(id).cloned())
                    .collect()
            })
            .unwrap_or_default()
    }
    
    /// Get all relations connected to an entity (both directions)
    pub fn relations_of(&self, entity_id: &str) -> Vec<Relation> {
        let mut result = self.outgoing_from(entity_id);
        result.extend(self.incoming_to(entity_id));
        result
    }
    
    /// Get neighboring entities (connected by any relation)
    pub fn neighbors(&self, entity_id: &str) -> Vec<Entity> {
        let relations = self.relations_of(entity_id);
        let entities = self.entities.read();
        
        let mut neighbor_ids: HashSet<String> = HashSet::new();
        
        for relation in relations {
            if relation.source == entity_id {
                neighbor_ids.insert(relation.target);
            } else {
                neighbor_ids.insert(relation.source);
            }
        }
        
        neighbor_ids.into_iter()
            .filter_map(|id| entities.get(&id).cloned())
            .collect()
    }
    
    /// Find relations of a specific type from an entity
    pub fn outgoing_of_type(&self, entity_id: &str, relation_type: &str) -> Vec<Relation> {
        self.outgoing_from(entity_id)
            .into_iter()
            .filter(|r| r.relation_type == relation_type)
            .collect()
    }
    
    /// BFS traversal from an entity up to max_hops
    pub fn traverse_bfs(
        &self,
        start_entity_id: &str,
        max_hops: usize,
        relation_types: Option<&[&str]>,
    ) -> Vec<(Entity, usize)> {
        let entities = self.entities.read();
        let relations = self.relations.read();
        let outgoing = self.outgoing_relations.read();
        
        let mut visited: HashSet<String> = HashSet::new();
        let mut result: Vec<(Entity, usize)> = Vec::new();
        let mut queue: Vec<(String, usize)> = vec![(start_entity_id.to_string(), 0)];
        
        visited.insert(start_entity_id.to_string());
        
        while let Some((current_id, depth)) = queue.pop() {
            if let Some(entity) = entities.get(&current_id) {
                if current_id != start_entity_id {
                    result.push((entity.clone(), depth));
                }
            }
            
            if depth >= max_hops {
                continue;
            }
            
            // Get outgoing relations
            if let Some(relation_ids) = outgoing.get(&current_id) {
                for rel_id in relation_ids {
                    if let Some(relation) = relations.get(rel_id) {
                        // Filter by relation type if specified
                        if let Some(types) = relation_types {
                            if !types.contains(&relation.relation_type.as_str()) {
                                continue;
                            }
                        }
                        
                        let target = &relation.target;
                        if !visited.contains(target) {
                            visited.insert(target.clone());
                            queue.push((target.clone(), depth + 1));
                        }
                    }
                }
            }
        }
        
        result
    }
    
    /// Find shortest path between two entities
    pub fn shortest_path(&self, start_id: &str, end_id: &str, max_hops: usize) -> Option<Vec<String>> {
        let outgoing = self.outgoing_relations.read();
        let relations = self.relations.read();
        
        let mut visited: HashSet<String> = HashSet::new();
        let mut queue: Vec<(String, Vec<String>)> = vec![(start_id.to_string(), vec![start_id.to_string()])];
        
        visited.insert(start_id.to_string());
        
        while let Some((current_id, path)) = queue.pop() {
            if current_id == end_id {
                return Some(path);
            }
            
            if path.len() > max_hops {
                continue;
            }
            
            if let Some(relation_ids) = outgoing.get(&current_id) {
                for rel_id in relation_ids {
                    if let Some(relation) = relations.get(rel_id) {
                        let target = &relation.target;
                        if !visited.contains(target) {
                            visited.insert(target.clone());
                            let mut new_path = path.clone();
                            new_path.push(target.clone());
                            queue.push((target.clone(), new_path));
                        }
                    }
                }
            }
        }
        
        None
    }
    
    /// Check if store has unsaved changes
    pub fn is_dirty(&self) -> bool {
        *self.dirty.read()
    }
    
    /// Flush changes to disk
    pub fn flush(&self) -> io::Result<()> {
        if self.is_dirty() {
            self.save()?;
        }
        Ok(())
    }
}

impl Drop for GraphStore {
    fn drop(&mut self) {
        let _ = self.flush();
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::tempdir;
    use crate::types::{entity_types, relation_types};

    #[test]
    fn test_graph_store_entities() {
        let dir = tempdir().unwrap();
        let store = GraphStore::open_or_create(dir.path()).unwrap();
        
        let entity = Entity::new("Albert Einstein", entity_types::PERSON)
            .with_property("birth_year".to_string(), serde_json::json!(1879));
        
        store.put_entity(entity.clone());
        
        assert_eq!(store.entity_count(), 1);
        assert!(store.contains_entity(&entity.id));
        
        let retrieved = store.get_entity(&entity.id).unwrap();
        assert_eq!(retrieved.name, "Albert Einstein");
    }

    #[test]
    fn test_graph_store_relations() {
        let dir = tempdir().unwrap();
        let store = GraphStore::open_or_create(dir.path()).unwrap();
        
        let e1 = Entity::new("Physics", entity_types::CONCEPT);
        let e2 = Entity::new("Mathematics", entity_types::CONCEPT);
        
        store.put_entity(e1.clone());
        store.put_entity(e2.clone());
        
        let relation = Relation::new(&e1.id, &e2.id, relation_types::RELATES_TO);
        store.put_relation(relation.clone());
        
        assert_eq!(store.relation_count(), 1);
        
        let outgoing = store.outgoing_from(&e1.id);
        assert_eq!(outgoing.len(), 1);
        assert_eq!(outgoing[0].target, e2.id);
    }

    #[test]
    fn test_graph_traversal() {
        let dir = tempdir().unwrap();
        let store = GraphStore::open_or_create(dir.path()).unwrap();
        
        // Create a simple graph: A -> B -> C
        let a = Entity::new("A", "node");
        let b = Entity::new("B", "node");
        let c = Entity::new("C", "node");
        
        store.put_entity(a.clone());
        store.put_entity(b.clone());
        store.put_entity(c.clone());
        
        store.put_relation(Relation::new(&a.id, &b.id, "connects"));
        store.put_relation(Relation::new(&b.id, &c.id, "connects"));
        
        // Traverse from A with max 2 hops
        let result = store.traverse_bfs(&a.id, 2, None);
        assert_eq!(result.len(), 2); // Should find B and C
        
        // Find path from A to C
        let path = store.shortest_path(&a.id, &c.id, 5).unwrap();
        assert_eq!(path.len(), 3); // A -> B -> C
    }

    #[test]
    fn test_graph_persistence() {
        let dir = tempdir().unwrap();
        
        let entity_id;
        
        // Create and save
        {
            let store = GraphStore::open_or_create(dir.path()).unwrap();
            let entity = Entity::new("Persistent", entity_types::CONCEPT);
            entity_id = entity.id.clone();
            store.put_entity(entity);
            store.save().unwrap();
        }
        
        // Reopen and verify
        {
            let store = GraphStore::open_or_create(dir.path()).unwrap();
            assert_eq!(store.entity_count(), 1);
            
            let entity = store.get_entity(&entity_id).unwrap();
            assert_eq!(entity.name, "Persistent");
        }
    }
}
