#!/usr/bin/env python3
"""
GOFAI Knowledge Graph for Hermes Fleet
Property graph for wizard lineage, tasks, and relationships
"""

import json
import hashlib
from datetime import datetime
from typing import Dict, List, Optional, Set, Any, Tuple
from dataclasses import dataclass, asdict
from pathlib import Path

@dataclass
class Node:
    """Graph node representing an entity"""
    id: str
    label: str  # Type: Wizard, Task, Artifact, etc.
    properties: Dict[str, Any]
    created_at: str = None
    
    def __post_init__(self):
        if self.created_at is None:
            self.created_at = datetime.now().isoformat()
    
    def to_dict(self) -> Dict:
        return {
            'id': self.id,
            'label': self.label,
            'properties': self.properties,
            'created_at': self.created_at
        }
    
    @classmethod
    def from_dict(cls, data: Dict) -> 'Node':
        return cls(**data)

@dataclass
class Edge:
    """Graph edge representing a relationship"""
    id: str
    source: str  # Source node ID
    target: str  # Target node ID
    relation: str  # Type of relationship
    properties: Dict[str, Any]
    created_at: str = None
    
    def __post_init__(self):
        if self.created_at is None:
            self.created_at = datetime.now().isoformat()
    
    def to_dict(self) -> Dict:
        return {
            'id': self.id,
            'source': self.source,
            'target': self.target,
            'relation': self.relation,
            'properties': self.properties,
            'created_at': self.created_at
        }

class KnowledgeGraph:
    """
    Property graph for fleet knowledge
    Stores wizards, tasks, artifacts and their relationships
    """
    
    def __init__(self, storage_path: Optional[str] = None):
        self.nodes: Dict[str, Node] = {}
        self.edges: Dict[str, Edge] = {}
        self.indexes: Dict[str, Dict[str, Set[str]]] = {
            'label': {},  # label -> node_ids
            'relation': {},  # relation -> edge_ids
        }
        self.storage_path = storage_path
        
        if storage_path and Path(storage_path).exists():
            self.load()
    
    def _generate_id(self, data: str) -> str:
        """Generate unique ID from data"""
        return hashlib.sha256(data.encode()).hexdigest()[:16]
    
    def add_node(self, label: str, properties: Dict, node_id: Optional[str] = None) -> str:
        """Add a node to the graph"""
        if node_id is None:
            node_id = self._generate_id(f"{label}:{json.dumps(properties, sort_keys=True)}")
        
        node = Node(id=node_id, label=label, properties=properties)
        self.nodes[node_id] = node
        
        # Update index
        if label not in self.indexes['label']:
            self.indexes['label'][label] = set()
        self.indexes['label'][label].add(node_id)
        
        return node_id
    
    def add_edge(self, source: str, target: str, relation: str, 
                 properties: Optional[Dict] = None, edge_id: Optional[str] = None) -> str:
        """Add an edge between nodes"""
        if properties is None:
            properties = {}
        
        if edge_id is None:
            edge_id = self._generate_id(f"{source}:{relation}:{target}")
        
        edge = Edge(id=edge_id, source=source, target=target, 
                   relation=relation, properties=properties)
        self.edges[edge_id] = edge
        
        # Update index
        if relation not in self.indexes['relation']:
            self.indexes['relation'][relation] = set()
        self.indexes['relation'][relation].add(edge_id)
        
        return edge_id
    
    def get_node(self, node_id: str) -> Optional[Node]:
        """Get a node by ID"""
        return self.nodes.get(node_id)
    
    def get_neighbors(self, node_id: str, relation: Optional[str] = None,
                     direction: str = 'out') -> List[Tuple[Node, Edge]]:
        """Get neighboring nodes"""
        results = []
        
        for edge in self.edges.values():
            if direction in ('out', 'both') and edge.source == node_id:
                if relation is None or edge.relation == relation:
                    target = self.nodes.get(edge.target)
                    if target:
                        results.append((target, edge))
            
            if direction in ('in', 'both') and edge.target == node_id:
                if relation is None or edge.relation == relation:
                    source = self.nodes.get(edge.source)
                    if source:
                        results.append((source, edge))
        
        return results
    
    def query(self, label: Optional[str] = None, 
              properties: Optional[Dict] = None) -> List[Node]:
        """Query nodes by label and/or properties"""
        results = []
        
        # Start with label filter if provided
        if label and label in self.indexes['label']:
            candidates = [self.nodes[nid] for nid in self.indexes['label'][label]]
        else:
            candidates = list(self.nodes.values())
        
        # Filter by properties
        if properties:
            for node in candidates:
                matches = all(
                    node.properties.get(k) == v 
                    for k, v in properties.items()
                )
                if matches:
                    results.append(node)
        else:
            results = candidates
        
        return results
    
    def find_path(self, start_id: str, end_id: str, 
                  max_depth: int = 5) -> Optional[List[Edge]]:
        """Find path between two nodes (BFS)"""
        if start_id not in self.nodes or end_id not in self.nodes:
            return None
        
        # BFS
        visited = {start_id}
        queue = [(start_id, [])]
        
        while queue:
            current_id, path = queue.pop(0)
            
            if current_id == end_id:
                return path
            
            if len(path) >= max_depth:
                continue
            
            # Get outgoing edges
            for edge in self.edges.values():
                if edge.source == current_id and edge.target not in visited:
                    visited.add(edge.target)
                    queue.append((edge.target, path + [edge]))
        
        return None
    
    def get_lineage(self, wizard_id: str) -> Dict:
        """Get full lineage tree for a wizard"""
        wizard = self.nodes.get(wizard_id)
        if not wizard or wizard.label != 'Wizard':
            return {}
        
        tree = {
            'wizard': wizard.properties.get('name', wizard_id),
            'father': None,
            'grandfather': None,
            'children': [],
            'siblings': []
        }
        
        # Find father
        father_edges = self.get_neighbors(wizard_id, relation='child_of', direction='out')
        for father_node, edge in father_edges:
            tree['father'] = father_node.properties.get('name', father_node.id)
            
            # Find grandfather
            grandfather_edges = self.get_neighbors(father_node.id, relation='child_of', direction='out')
            for grandpa_node, _ in grandfather_edges:
                tree['grandfather'] = grandpa_node.properties.get('name', grandpa_node.id)
            
            # Find siblings (other children of father)
            sibling_edges = self.get_neighbors(father_node.id, relation='child_of', direction='in')
            for sibling_node, _ in sibling_edges:
                if sibling_node.id != wizard_id:
                    tree['siblings'].append(sibling_node.properties.get('name', sibling_node.id))
        
        # Find children
        child_edges = self.get_neighbors(wizard_id, relation='child_of', direction='in')
        for child_node, _ in child_edges:
            tree['children'].append(child_node.properties.get('name', child_node.id))
        
        return tree
    
    def get_task_dependencies(self, task_id: str) -> Dict:
        """Get task dependency graph"""
        task = self.nodes.get(task_id)
        if not task:
            return {}
        
        deps = {
            'task': task.properties.get('name', task_id),
            'depends_on': [],
            'blocks': [],
            'assignee': None
        }
        
        # Find dependencies
        dep_edges = self.get_neighbors(task_id, relation='depends_on', direction='out')
        for dep_task, _ in dep_edges:
            deps['depends_on'].append(dep_task.properties.get('name', dep_task.id))
        
        # Find what this blocks
        blocker_edges = self.get_neighbors(task_id, relation='depends_on', direction='in')
        for blocked_task, _ in blocker_edges:
            deps['blocks'].append(blocked_task.properties.get('name', blocked_task.id))
        
        # Find assignee
        assignee_edges = self.get_neighbors(task_id, relation='assigned_to', direction='out')
        for wizard, _ in assignee_edges:
            deps['assignee'] = wizard.properties.get('name', wizard.id)
        
        return deps
    
    def save(self) -> None:
        """Save graph to disk"""
        if not self.storage_path:
            return
        
        data = {
            'nodes': {nid: node.to_dict() for nid, node in self.nodes.items()},
            'edges': {eid: edge.to_dict() for eid, edge in self.edges.items()},
            'indexes': {k: {ik: list(iv) for ik, iv in v.items()} 
                       for k, v in self.indexes.items()}
        }
        
        Path(self.storage_path).write_text(json.dumps(data, indent=2))
    
    def load(self) -> None:
        """Load graph from disk"""
        if not self.storage_path or not Path(self.storage_path).exists():
            return
        
        data = json.loads(Path(self.storage_path).read_text())
        
        # Load nodes
        for nid, node_data in data.get('nodes', {}).items():
            self.nodes[nid] = Node.from_dict(node_data)
        
        # Load edges
        for eid, edge_data in data.get('edges', {}).items():
            self.edges[eid] = Edge(**edge_data)
        
        # Load indexes
        self.indexes = {
            k: {ik: set(iv) for ik, iv in v.items()}
            for k, v in data.get('indexes', {}).items()
        }
    
    def export_graphviz(self) -> str:
        """Export as Graphviz DOT format"""
        lines = ['digraph Fleet {', '  rankdir=TB;']
        
        # Add nodes
        for node in self.nodes.values():
            label = node.properties.get('name', node.id[:8])
            shape = 'box' if node.label == 'Wizard' else 'ellipse'
            color = 'lightblue' if node.label == 'Wizard' else 'lightgreen'
            lines.append(f'  "{node.id}" [label="{label}", shape={shape}, style=filled, fillcolor={color}];')
        
        # Add edges
        for edge in self.edges.values():
            lines.append(f'  "{edge.source}" -> "{edge.target}" [label="{edge.relation}"];')
        
        lines.append('}')
        return '\n'.join(lines)
    
    def get_stats(self) -> Dict:
        """Get graph statistics"""
        return {
            'nodes': len(self.nodes),
            'edges': len(self.edges),
            'by_label': {label: len(nodes) for label, nodes in self.indexes['label'].items()},
            'by_relation': {rel: len(edges) for rel, edges in self.indexes['relation'].items()}
        }

class FleetKnowledgeBase:
    """
    High-level interface for fleet knowledge operations
    """
    
    def __init__(self, storage_path: str = '/root/wizards/allegro/gofai/fleet_kg.json'):
        self.kg = KnowledgeGraph(storage_path)
    
    def register_wizard(self, name: str, role: str, father: Optional[str] = None,
                       api_port: int = 0, model: str = "", capabilities: List[str] = None) -> str:
        """Register a wizard in the fleet"""
        
        properties = {
            'name': name,
            'role': role,
            'father': father,
            'api_port': api_port,
            'model': model,
            'capabilities': capabilities or [],
            'status': 'active'
        }
        
        wizard_id = self.kg.add_node('Wizard', properties, node_id=f'wizard_{name.lower()}')
        
        # Add father relationship if specified
        if father:
            father_id = f'wizard_{father.lower()}'
            if father_id in self.kg.nodes:
                self.kg.add_edge(wizard_id, father_id, 'child_of')
        
        self.kg.save()
        return wizard_id
    
    def register_task(self, name: str, priority: str, assignee: str,
                     dependencies: List[str] = None) -> str:
        """Register a task"""
        
        properties = {
            'name': name,
            'priority': priority,
            'status': 'pending',
            'assignee': assignee
        }
        
        task_id = self.kg.add_node('Task', properties)
        
        # Link to assignee
        wizard_id = f'wizard_{assignee.lower()}'
        if wizard_id in self.kg.nodes:
            self.kg.add_edge(task_id, wizard_id, 'assigned_to')
        
        # Add dependencies
        if dependencies:
            for dep_name in dependencies:
                # Find task by name
                for nid, node in self.kg.nodes.items():
                    if node.label == 'Task' and node.properties.get('name') == dep_name:
                        self.kg.add_edge(task_id, nid, 'depends_on')
                        break
        
        self.kg.save()
        return task_id
    
    def get_wizard_info(self, name: str) -> Dict:
        """Get comprehensive info about a wizard"""
        wizard_id = f'wizard_{name.lower()}'
        
        lineage = self.kg.get_lineage(wizard_id)
        
        # Get assigned tasks
        tasks = []
        for node in self.kg.query(label='Task'):
            if node.properties.get('assignee') == name:
                tasks.append({
                    'name': node.properties.get('name'),
                    'priority': node.properties.get('priority'),
                    'status': node.properties.get('status')
                })
        
        wizard = self.kg.get_node(wizard_id)
        
        return {
            'name': name,
            'lineage': lineage,
            'tasks': tasks,
            'capabilities': wizard.properties.get('capabilities', []) if wizard else [],
            'api_port': wizard.properties.get('api_port', 0) if wizard else 0
        }
    
    def who_can_do(self, capability: str) -> List[str]:
        """Find wizards with a specific capability"""
        results = []
        for node in self.kg.query(label='Wizard'):
            caps = node.properties.get('capabilities', [])
            if capability in caps:
                results.append(node.properties.get('name', node.id))
        return results