"""
Unified Graph Builder.

Aggregates the raw, multi-edge provenance graph (GraphStore) into a
clean, conflict-resolved knowledge graph (Unity Graph).

Logic:
- Consolidates multiple edges of same type between A and B into one edge.
- Aggregates trust scores (e.g. 1.0 + 1.0 => higher confidence).
- Merges properties using consensus or latest-wins strategy.
"""
from __future__ import annotations
import networkx as nx
from typing import Optional, Any
from collections import defaultdict

# Definition of the source graph interface (GraphStore)
# We don't import GraphStore directly to avoid circular imports if it's in persistence
# We just need an object with _G (MultiDiGraph)


class UnifiedGraphBuilder:
    """
    Builds a unified knowledge graph from a provenance graph store.
    """

    def __init__(self, graph_store):
        """
        Initialize with a GraphStore instance.

        Args:
            graph_store: Object with a ._G (nx.MultiDiGraph) attribute.
        """
        self.store = graph_store

    def build(self, namespace: Optional[str] = None) -> nx.DiGraph:
        """
        Build the unified graph.

        Args:
            namespace: Optional namespace to filter by.

        Returns:
            nx.DiGraph: A simple directed graph with consolidated facts.
        """
        if not hasattr(self.store, "_G"):
            return nx.DiGraph()

        raw_graph: nx.MultiDiGraph = self.store._G
        unified = nx.DiGraph()

        # 1. Collect Edges
        # Map: (u, v, relation_type) -> list of edge_data
        grouped_edges = defaultdict(list)

        for u, v, key, data in raw_graph.edges(data=True, keys=True):
            # Filter by namespace
            if namespace and data.get("namespace") != namespace:
                continue

            # Check node namespaces if applicable
            if namespace:
                if namespace not in raw_graph.nodes[u].get("namespaces", set()):
                    continue
                if namespace not in raw_graph.nodes[v].get("namespaces", set()):
                    continue

            # Relation normalization
            relation = data.get("relation_type") or data.get("relation") or "RELATED_TO"
            relation = str(relation).strip().lower()

            grouped_edges[(u, v, relation)].append(data)

        # 2. Add Nodes
        # We need to add all nodes involved in selected edges
        nodes_to_add = set()
        for (u, v, _), _ in grouped_edges.items():
            nodes_to_add.add(u)
            nodes_to_add.add(v)
        for node in nodes_to_add:
            # Copy node attributes from raw graph
            attrs = dict(raw_graph.nodes[node])
            unified.add_node(node, **attrs)

        # 3. Merge Edges
        for (u, v, relation), edges_data in grouped_edges.items():
            consensus_data = self._merge_edge_group(edges_data)
            consensus_data["weight"] = self._calculate_weight(edges_data)
            consensus_data["source_count"] = len(edges_data)

            # Add to unified graph
            unified.add_edge(u, v, relation=relation, **consensus_data)

        return unified

    def _merge_edge_group(self, edges_data: list[dict[str, Any]]) -> dict[str, Any]:
        """
        Merge multiple edge attributes into one consensus set.
        """
        # Start with the stats from the most trusted/recent edge
        # For now, simple strategy: merge all non-conflicting keys
        merged = {}

        # Calculate aggregate trust
        total_trust = 0.0

        # For property conflict resolution (consensus)
        property_counts = defaultdict(lambda: defaultdict(float))

        for data in edges_data:
            trust = float(data.get("trust", 1.0))
            total_trust += trust

            for k, v in data.items():
                if k in ("trust", "weight", "frequency", "episode", "namespace", "relation", "relation_type"):
                    continue
                # Count votes for each property value weighted by trust
                property_counts[k][v] += trust

        # Resolve properties
        for k, value_counts in property_counts.items():
            # Pick value with highest weighted vote
            best_val = max(value_counts.items(), key=lambda x: x[1])[0]
            merged[k] = best_val

        return merged

    def _calculate_weight(self, edges_data: list[dict[str, Any]]) -> float:
        """
        Calculate edge weight for traversal algorithms.
        More confirmation = lower cost (if distance) or higher capacity (if flow).
        Here we return 'strength' (higher is better).
        """
        # Sum of trusts
        return sum(float(d.get("trust", 1.0)) for d in edges_data)
