"""Memory Graph - Cross-Workflow Knowledge Base

A knowledge graph that connects findings across all workflows,
enabling intelligent correlation and learning.

Features:
- Add findings from any workflow as nodes
- Connect related findings with typed edges
- Query for similar past findings
- Traverse relationships to find root causes

Storage: JSON file in patterns/ directory

Copyright 2025 Smart AI Memory, LLC
Licensed under Fair Source 0.9
"""

import hashlib
import json
from collections import defaultdict, deque
from datetime import datetime
from pathlib import Path
from typing import Any

from attune.config import _validate_file_path

from .edges import REVERSE_EDGE_TYPES, Edge, EdgeType
from .nodes import Node, NodeType


class MemoryGraph:
    """Knowledge graph for cross-workflow intelligence.

    Stores nodes (findings) and edges (relationships) discovered
    by workflows, enabling pattern correlation across sessions.

    Usage:
        graph = MemoryGraph()

        # Add a bug finding
        bug_id = graph.add_finding(
            workflow="bug-predict",
            finding={
                "type": "bug",
                "name": "Null reference in auth.py",
                "description": "Missing null check on user object",
                "file": "src/auth.py",
                "line": 42,
                "severity": "high"
            }
        )

        # Connect to a fix
        fix_id = graph.add_finding(
            workflow="bug-predict",
            finding={
                "type": "fix",
                "name": "Add null check",
                "description": "Added guard clause for user object"
            }
        )
        graph.add_edge(bug_id, fix_id, EdgeType.FIXED_BY)

        # Find similar bugs
        similar = graph.find_similar({"name": "Null reference"})
    """

    def __init__(self, path: str | Path = "patterns/memory_graph.json"):
        """Initialize the memory graph.

        Args:
            path: Path to JSON storage file

        """
        self.path = Path(path)
        self.nodes: dict[str, Node] = {}
        self.edges: list[Edge] = []

        # Indexes for fast lookup
        self._edges_by_source: dict[str, list[Edge]] = defaultdict(list)
        self._edges_by_target: dict[str, list[Edge]] = defaultdict(list)
        self._nodes_by_type: dict[NodeType, list[str]] = defaultdict(list)
        self._nodes_by_workflow: dict[str, list[str]] = defaultdict(list)
        self._nodes_by_file: dict[str, list[str]] = defaultdict(list)

        self._load()

    def _load(self) -> None:
        """Load graph from JSON file."""
        if not self.path.exists():
            # Ensure directory exists
            self.path.parent.mkdir(parents=True, exist_ok=True)
            self._save()
            return

        try:
            with open(self.path) as f:
                data = json.load(f)

            # Load nodes
            for node_data in data.get("nodes", []):
                node = Node.from_dict(node_data)
                self.nodes[node.id] = node
                self._index_node(node)

            # Load edges
            for edge_data in data.get("edges", []):
                edge = Edge.from_dict(edge_data)
                self.edges.append(edge)
                self._index_edge(edge)

        except (json.JSONDecodeError, KeyError) as e:
            print(f"Warning: Could not load graph from {self.path}: {e}")
            self.nodes = {}
            self.edges = []

    def _save(self) -> None:
        """Save graph to JSON file."""
        data = {
            "version": "1.0",
            "updated_at": datetime.now().isoformat(),
            "node_count": len(self.nodes),
            "edge_count": len(self.edges),
            "nodes": [node.to_dict() for node in self.nodes.values()],
            "edges": [edge.to_dict() for edge in self.edges],
        }

        self.path.parent.mkdir(parents=True, exist_ok=True)
        validated_path = _validate_file_path(str(self.path))
        with open(validated_path, "w") as f:
            json.dump(data, f, indent=2)

    def _index_node(self, node: Node) -> None:
        """Add node to indexes."""
        self._nodes_by_type[node.type].append(node.id)
        if node.source_workflow:
            self._nodes_by_workflow[node.source_workflow].append(node.id)
        if node.source_file:
            self._nodes_by_file[node.source_file].append(node.id)

    def _index_edge(self, edge: Edge) -> None:
        """Add edge to indexes."""
        self._edges_by_source[edge.source_id].append(edge)
        self._edges_by_target[edge.target_id].append(edge)

    def _generate_id(self, finding: dict[str, Any]) -> str:
        """Generate unique ID for a finding."""
        # Create hash from content
        content = json.dumps(finding, sort_keys=True)
        hash_val = hashlib.sha256(content.encode()).hexdigest()[:12]
        timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
        return f"{finding.get('type', 'node')}_{timestamp}_{hash_val}"

    def add_finding(self, workflow: str, finding: dict[str, Any]) -> str:
        """Add a finding from any workflow, return node ID.

        Args:
            workflow: Name of the workflow adding this finding
            finding: Dict with at least 'type' and 'name' keys

        Returns:
            Node ID for the created node

        Example:
            node_id = graph.add_finding(
                workflow="security-audit",
                finding={
                    "type": "vulnerability",
                    "name": "SQL Injection in query builder",
                    "description": "User input not sanitized",
                    "file": "src/db/query.py",
                    "line": 156,
                    "severity": "critical"
                }
            )

        """
        node_id = self._generate_id(finding)

        # Map finding type to NodeType
        type_str = finding.get("type", "pattern")
        try:
            node_type = NodeType(type_str)
        except ValueError:
            node_type = NodeType.PATTERN

        node = Node(
            id=node_id,
            type=node_type,
            name=finding.get("name", "Unnamed finding"),
            description=finding.get("description", ""),
            source_workflow=workflow,
            source_file=finding.get("file", finding.get("source_file", "")),
            source_line=finding.get("line", finding.get("source_line")),
            severity=finding.get("severity", ""),
            confidence=finding.get("confidence", 1.0),
            metadata=finding.get("metadata", {}),
            tags=finding.get("tags", []),
        )

        self.nodes[node_id] = node
        self._index_node(node)
        self._save()

        return node_id

    def add_edge(
        self,
        source_id: str,
        target_id: str,
        edge_type: EdgeType,
        description: str = "",
        workflow: str = "",
        weight: float = 1.0,
        bidirectional: bool = False,
    ) -> str:
        """Add an edge between two nodes.

        Args:
            source_id: Source node ID
            target_id: Target node ID
            edge_type: Type of relationship
            description: Optional description of the relationship
            workflow: Workflow that created this edge
            weight: Strength of relationship (0.0 - 1.0)
            bidirectional: If True, also create reverse edge

        Returns:
            Edge ID

        """
        if source_id not in self.nodes:
            raise ValueError(f"Source node not found: {source_id}")
        if target_id not in self.nodes:
            raise ValueError(f"Target node not found: {target_id}")

        edge = Edge(
            source_id=source_id,
            target_id=target_id,
            type=edge_type,
            description=description,
            source_workflow=workflow,
            weight=weight,
        )

        self.edges.append(edge)
        self._index_edge(edge)

        # Optionally create reverse edge
        if bidirectional and edge_type in REVERSE_EDGE_TYPES:
            reverse_type = REVERSE_EDGE_TYPES[edge_type]
            reverse_edge = Edge(
                source_id=target_id,
                target_id=source_id,
                type=reverse_type,
                description=description,
                source_workflow=workflow,
                weight=weight,
            )
            self.edges.append(reverse_edge)
            self._index_edge(reverse_edge)

        self._save()
        return edge.id

    def get_node(self, node_id: str) -> Node | None:
        """Get a node by ID."""
        return self.nodes.get(node_id)

    def find_related(
        self,
        node_id: str,
        edge_types: list[EdgeType] | None = None,
        direction: str = "outgoing",
        max_depth: int = 1,
    ) -> list[Node]:
        """Find related nodes via specified edge types.

        Args:
            node_id: Starting node ID
            edge_types: List of edge types to follow (None = all)
            direction: "outgoing", "incoming", or "both"
            max_depth: Maximum traversal depth

        Returns:
            List of related nodes

        """
        if node_id not in self.nodes:
            return []

        visited: set[str] = {node_id}
        result: list[Node] = []
        current_level: set[str] = {node_id}

        for _ in range(max_depth):
            next_level: set[str] = set()

            for current_id in current_level:
                edges_to_check: list[Edge] = []

                if direction in ("outgoing", "both"):
                    edges_to_check.extend(self._edges_by_source.get(current_id, []))
                if direction in ("incoming", "both"):
                    edges_to_check.extend(self._edges_by_target.get(current_id, []))

                for edge in edges_to_check:
                    if edge_types and edge.type not in edge_types:
                        continue

                    # Get the other node
                    other_id = edge.target_id if edge.source_id == current_id else edge.source_id

                    if other_id not in visited:
                        visited.add(other_id)
                        next_level.add(other_id)
                        if other_id in self.nodes:
                            result.append(self.nodes[other_id])

            if not next_level:
                break
            current_level = next_level

        return result

    def find_similar(
        self,
        finding: dict[str, Any],
        threshold: float = 0.5,
        limit: int = 10,
    ) -> list[tuple[Node, float]]:
        """Find similar past findings.

        Uses simple text similarity on name and description.

        Args:
            finding: Dict with 'name' and/or 'description'
            threshold: Minimum similarity score (0.0 - 1.0)
            limit: Maximum results to return

        Returns:
            List of (node, similarity_score) tuples

        """
        query_name = finding.get("name", "").lower()
        query_desc = finding.get("description", "").lower()
        query_type = finding.get("type")
        query_file = finding.get("file", "")

        results: list[tuple[Node, float]] = []

        for node in self.nodes.values():
            score = 0.0
            factors = 0.0

            # Name similarity (word overlap)
            if query_name and node.name:
                name_words = set(query_name.split())
                node_words = set(node.name.lower().split())
                if name_words and node_words:
                    overlap = len(name_words & node_words)
                    union = len(name_words | node_words)
                    score += (overlap / union) * 0.5
                    factors += 0.5

            # Description similarity
            if query_desc and node.description:
                desc_words = set(query_desc.split())
                node_words = set(node.description.lower().split())
                if desc_words and node_words:
                    overlap = len(desc_words & node_words)
                    union = len(desc_words | node_words)
                    score += (overlap / union) * 0.3
                    factors += 0.3

            # Type match bonus
            if query_type:
                try:
                    if node.type == NodeType(query_type):
                        score += 0.15
                        factors += 0.15
                except ValueError:
                    pass

            # File match bonus
            if query_file and node.source_file:
                if query_file == node.source_file:
                    score += 0.05
                    factors += 0.05

            # Normalize
            if factors > 0:
                score = score / factors

            if score >= threshold:
                results.append((node, score))

        # Sort by score descending
        results.sort(key=lambda x: x[1], reverse=True)
        return results[:limit]

    def find_by_type(self, node_type: NodeType) -> list[Node]:
        """Find all nodes of a specific type."""
        node_ids = self._nodes_by_type.get(node_type, [])
        return [self.nodes[nid] for nid in node_ids if nid in self.nodes]

    def find_by_workflow(self, workflow: str) -> list[Node]:
        """Find all nodes created by a specific workflow."""
        node_ids = self._nodes_by_workflow.get(workflow, [])
        return [self.nodes[nid] for nid in node_ids if nid in self.nodes]

    def find_by_file(self, file_path: str) -> list[Node]:
        """Find all nodes related to a specific file."""
        node_ids = self._nodes_by_file.get(file_path, [])
        return [self.nodes[nid] for nid in node_ids if nid in self.nodes]

    def get_path(
        self,
        source_id: str,
        target_id: str,
        edge_types: list[EdgeType] | None = None,
        max_depth: int = 5,
    ) -> list[tuple[Node, Edge | None]]:
        """Find a path between two nodes using BFS.

        Args:
            source_id: Starting node ID
            target_id: Target node ID
            edge_types: Edge types to traverse (None = all)
            max_depth: Maximum path length

        Returns:
            List of (node, edge_to_node) tuples representing the path

        """
        if source_id not in self.nodes or target_id not in self.nodes:
            return []

        # BFS with path tracking (deque for O(1) popleft)
        visited: set[str] = {source_id}
        queue: deque[list[tuple[str, Edge | None]]] = deque([[(source_id, None)]])

        while queue:
            path = queue.popleft()
            current_id = path[-1][0]

            if len(path) > max_depth:
                continue

            if current_id == target_id:
                return [(self.nodes[nid], edge) for nid, edge in path]

            for edge in self._edges_by_source.get(current_id, []):
                if edge_types and edge.type not in edge_types:
                    continue
                if edge.target_id not in visited:
                    visited.add(edge.target_id)
                    new_path = path + [(edge.target_id, edge)]
                    queue.append(new_path)

        return []

    def get_statistics(self) -> dict[str, Any]:
        """Get graph statistics."""
        type_counts: dict[str, int] = defaultdict(int)
        workflow_counts: dict[str, int] = defaultdict(int)
        severity_counts: dict[str, int] = defaultdict(int)

        for node in self.nodes.values():
            type_counts[node.type.value] += 1
            if node.source_workflow:
                workflow_counts[node.source_workflow] += 1
            if node.severity:
                severity_counts[node.severity] += 1

        edge_type_counts: dict[str, int] = defaultdict(int)
        for edge in self.edges:
            edge_type_counts[edge.type.value] += 1

        return {
            "total_nodes": len(self.nodes),
            "total_edges": len(self.edges),
            "nodes_by_type": dict(type_counts),
            "nodes_by_workflow": dict(workflow_counts),
            "nodes_by_severity": dict(severity_counts),
            "edges_by_type": dict(edge_type_counts),
            "unique_files": len(self._nodes_by_file),
        }

    def update_node(self, node_id: str, updates: dict[str, Any]) -> bool:
        """Update a node's properties.

        Args:
            node_id: Node to update
            updates: Dict of properties to update

        Returns:
            True if updated, False if node not found

        """
        if node_id not in self.nodes:
            return False

        node = self.nodes[node_id]

        if "status" in updates:
            node.status = updates["status"]
        if "description" in updates:
            node.description = updates["description"]
        if "severity" in updates:
            node.severity = updates["severity"]
        if "tags" in updates:
            node.tags = updates["tags"]
        if "metadata" in updates:
            node.metadata.update(updates["metadata"])

        node.updated_at = datetime.now()
        self._save()
        return True

    def delete_node(self, node_id: str) -> bool:
        """Delete a node and its connected edges.

        Args:
            node_id: Node to delete

        Returns:
            True if deleted, False if not found

        """
        if node_id not in self.nodes:
            return False

        # Remove from indexes
        node = self.nodes[node_id]
        if node.type in self._nodes_by_type:
            self._nodes_by_type[node.type] = [
                nid for nid in self._nodes_by_type[node.type] if nid != node_id
            ]
        if node.source_workflow in self._nodes_by_workflow:
            self._nodes_by_workflow[node.source_workflow] = [
                nid for nid in self._nodes_by_workflow[node.source_workflow] if nid != node_id
            ]
        if node.source_file in self._nodes_by_file:
            self._nodes_by_file[node.source_file] = [
                nid for nid in self._nodes_by_file[node.source_file] if nid != node_id
            ]

        # Remove connected edges
        self.edges = [e for e in self.edges if e.source_id != node_id and e.target_id != node_id]
        if node_id in self._edges_by_source:
            del self._edges_by_source[node_id]
        if node_id in self._edges_by_target:
            del self._edges_by_target[node_id]

        # Remove node
        del self.nodes[node_id]
        self._save()
        return True

    def clear(self) -> None:
        """Clear all nodes and edges."""
        self.nodes = {}
        self.edges = []
        self._edges_by_source = defaultdict(list)
        self._edges_by_target = defaultdict(list)
        self._nodes_by_type = defaultdict(list)
        self._nodes_by_workflow = defaultdict(list)
        self._nodes_by_file = defaultdict(list)
        self._save()
