"""
Counterfactual Edges
====================

Provides counterfactual relationship tracking between episodes.

Edge types:
- contradicts: Direct contradiction between episodes
- supersedes: Updated/newer information replaces older
- obsoletes: Information is no longer valid
- revises: Partial update/correction
- supports: Evidence supporting another episode
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from typing import Any, Dict, List, Optional, Set, Tuple

logger = logging.getLogger(__name__)


class CounterfactualEdgeType(Enum):
    """Types of counterfactual relationships between episodes."""
    CONTRADICTS = "contradicts"     # Direct contradiction
    SUPERSEDES = "supersedes"       # Updated information replaces older
    OBSOLETES = "obsoletes"         # Information no longer valid
    REVISES = "revises"             # Partial update/correction
    SUPPORTS = "supports"           # Evidence supporting another episode
    WEAKENS = "weakens"             # Evidence weakening another episode


@dataclass
class CounterfactualEdge:
    """
    A counterfactual relationship between two episodes.

    Direction: source_episode_id -> target_episode_id
    Example: (new_fact) -[supersedes]-> (old_fact)
    """
    source_episode_id: str
    target_episode_id: str
    edge_type: CounterfactualEdgeType
    confidence: float = 1.0
    reason: Optional[str] = None
    created_at: datetime = field(default_factory=datetime.utcnow)
    metadata: Dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for serialization."""
        return {
            "source_episode_id": self.source_episode_id,
            "target_episode_id": self.target_episode_id,
            "edge_type": self.edge_type.value,
            "confidence": self.confidence,
            "reason": self.reason,
            "created_at": self.created_at.isoformat(),
            "metadata": self.metadata,
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "CounterfactualEdge":
        """Create from dictionary."""
        return cls(
            source_episode_id=data["source_episode_id"],
            target_episode_id=data["target_episode_id"],
            edge_type=CounterfactualEdgeType(data["edge_type"]),
            confidence=data.get("confidence", 1.0),
            reason=data.get("reason"),
            created_at=datetime.fromisoformat(data["created_at"]) if "created_at" in data else datetime.utcnow(),
            metadata=data.get("metadata", {}),
        )

    def inverse(self) -> "CounterfactualEdge":
        """Get the inverse edge (if applicable)."""
        inverse_types = {
            CounterfactualEdgeType.CONTRADICTS: CounterfactualEdgeType.CONTRADICTS,
            CounterfactualEdgeType.SUPERSEDES: None,  # Not invertible
            CounterfactualEdgeType.SUPPORTS: CounterfactualEdgeType.SUPPORTS,
            CounterfactualEdgeType.WEAKENS: CounterfactualEdgeType.WEAKENS,
        }

        inverse_type = inverse_types.get(self.edge_type)
        if inverse_type is None:
            raise ValueError(f"Edge type {self.edge_type} is not invertible")

        return CounterfactualEdge(
            source_episode_id=self.target_episode_id,
            target_episode_id=self.source_episode_id,
            edge_type=inverse_type,
            confidence=self.confidence,
            reason=self.reason,
            metadata=self.metadata,
        )


@dataclass
class CounterfactualChain:
    """A chain of counterfactual relationships."""
    edges: List[CounterfactualEdge]
    start_episode_id: str
    end_episode_id: str

    @property
    def length(self) -> int:
        return len(self.edges)

    @property
    def total_confidence(self) -> float:
        """Multiply confidences along the chain."""
        if not self.edges:
            return 0.0
        conf = 1.0
        for edge in self.edges:
            conf *= edge.confidence
        return conf

    def to_dict(self) -> Dict[str, Any]:
        return {
            "edges": [e.to_dict() for e in self.edges],
            "start_episode_id": self.start_episode_id,
            "end_episode_id": self.end_episode_id,
            "length": self.length,
            "total_confidence": self.total_confidence,
        }


class CounterfactualGraph:
    """
    Manages counterfactual relationships between episodes.

    Provides efficient queries for:
    - What supersedes this episode?
    - What does this episode contradict?
    - What is the latest version of this fact?
    """

    def __init__(self) -> None:
        """Initialize the counterfactual graph."""
        # Adjacency lists: episode_id -> list of outgoing edges
        self._outgoing: Dict[str, List[CounterfactualEdge]] = {}
        # Reverse adjacency: episode_id -> list of incoming edges
        self._incoming: Dict[str, List[CounterfactualEdge]] = {}
        # Index by edge type for fast queries
        self._by_type: Dict[CounterfactualEdgeType, List[CounterfactualEdge]] = {
            t: [] for t in CounterfactualEdgeType
        }

    def add_edge(self, edge: CounterfactualEdge) -> None:
        """Add a counterfactual edge."""
        # Add to outgoing
        if edge.source_episode_id not in self._outgoing:
            self._outgoing[edge.source_episode_id] = []
        self._outgoing[edge.source_episode_id].append(edge)

        # Add to incoming
        if edge.target_episode_id not in self._incoming:
            self._incoming[edge.target_episode_id] = []
        self._incoming[edge.target_episode_id].append(edge)

        # Add to type index
        self._by_type[edge.edge_type].append(edge)

    def remove_edge(
        self,
        source_id: str,
        target_id: str,
        edge_type: Optional[CounterfactualEdgeType] = None,
    ) -> bool:
        """Remove an edge (optionally filtered by type)."""
        removed = False

        # Remove from outgoing
        if source_id in self._outgoing:
            original_len = len(self._outgoing[source_id])
            self._outgoing[source_id] = [
                e for e in self._outgoing[source_id]
                if not (e.target_episode_id == target_id
                        and (edge_type is None or e.edge_type == edge_type))
            ]
            removed = len(self._outgoing[source_id]) < original_len

        # Remove from incoming
        if target_id in self._incoming:
            self._incoming[target_id] = [
                e for e in self._incoming[target_id]
                if not (e.source_episode_id == source_id
                        and (edge_type is None or e.edge_type == edge_type))
            ]

        # Remove from type index
        for t in (edge_type,) if edge_type else CounterfactualEdgeType:
            self._by_type[t] = [
                e for e in self._by_type[t]
                if not (e.source_episode_id == source_id
                        and e.target_episode_id == target_id)
            ]

        return removed

    def remove_episode(self, episode_id: str) -> int:
        """Remove all edges involving an episode."""
        count = 0

        # Remove outgoing edges
        if episode_id in self._outgoing:
            for edge in self._outgoing[episode_id]:
                # Remove from target's incoming
                if edge.target_episode_id in self._incoming:
                    self._incoming[edge.target_episode_id] = [
                        e for e in self._incoming[edge.target_episode_id]
                        if e.source_episode_id != episode_id
                    ]
                # Remove from type index
                self._by_type[edge.edge_type] = [
                    e for e in self._by_type[edge.edge_type]
                    if e.source_episode_id != episode_id
                ]
                count += 1
            del self._outgoing[episode_id]

        # Remove incoming edges
        if episode_id in self._incoming:
            for edge in self._incoming[episode_id]:
                # Remove from source's outgoing
                if edge.source_episode_id in self._outgoing:
                    self._outgoing[edge.source_episode_id] = [
                        e for e in self._outgoing[edge.source_episode_id]
                        if e.target_episode_id != episode_id
                    ]
                # Remove from type index
                self._by_type[edge.edge_type] = [
                    e for e in self._by_type[edge.edge_type]
                    if e.target_episode_id != episode_id
                ]
                count += 1
            del self._incoming[episode_id]

        return count

    def get_outgoing(
        self,
        episode_id: str,
        edge_type: Optional[CounterfactualEdgeType] = None,
    ) -> List[CounterfactualEdge]:
        """Get outgoing edges from an episode."""
        edges = self._outgoing.get(episode_id, [])
        if edge_type:
            edges = [e for e in edges if e.edge_type == edge_type]
        return edges

    def get_incoming(
        self,
        episode_id: str,
        edge_type: Optional[CounterfactualEdgeType] = None,
    ) -> List[CounterfactualEdge]:
        """Get incoming edges to an episode."""
        edges = self._incoming.get(episode_id, [])
        if edge_type:
            edges = [e for e in edges if e.edge_type == edge_type]
        return edges

    def get_contradictions(self, episode_id: str) -> List[str]:
        """Get all episodes that contradict this one."""
        contradicting = set()

        # Outgoing contradicts edges
        for edge in self.get_outgoing(episode_id, CounterfactualEdgeType.CONTRADICTS):
            contradicting.add(edge.target_episode_id)

        # Incoming contradicts edges (symmetric)
        for edge in self.get_incoming(episode_id, CounterfactualEdgeType.CONTRADICTS):
            contradicting.add(edge.source_episode_id)

        return list(contradicting)

    def get_superseded_by(self, episode_id: str) -> List[str]:
        """Get episodes that supersede this one (newer versions)."""
        return [
            edge.source_episode_id
            for edge in self.get_incoming(episode_id, CounterfactualEdgeType.SUPERSEDES)
        ]

    def get_supersedes(self, episode_id: str) -> List[str]:
        """Get episodes that this one supersedes (older versions)."""
        return [
            edge.target_episode_id
            for edge in self.get_outgoing(episode_id, CounterfactualEdgeType.SUPERSEDES)
        ]

    def get_latest_version(self, episode_id: str, max_hops: int = 10) -> str:
        """
        Follow supersedes chain to find the latest version.

        Args:
            episode_id: Starting episode
            max_hops: Maximum chain length to follow

        Returns:
            ID of the latest version (or original if no supersedes)
        """
        current = episode_id
        visited: Set[str] = {current}

        for _ in range(max_hops):
            superseded_by = self.get_superseded_by(current)
            if not superseded_by:
                break

            # Take the most confident superseding edge
            edges = self.get_incoming(current, CounterfactualEdgeType.SUPERSEDES)
            best_edge = max(edges, key=lambda e: e.confidence)
            next_id = best_edge.source_episode_id

            if next_id in visited:
                logger.warning(f"Cycle detected in supersedes chain at {next_id}")
                break

            visited.add(next_id)
            current = next_id

        return current

    def get_original_version(self, episode_id: str, max_hops: int = 10) -> str:
        """
        Follow supersedes chain backward to find the original version.

        Args:
            episode_id: Starting episode
            max_hops: Maximum chain length to follow

        Returns:
            ID of the original version (or input if no superseded episodes)
        """
        current = episode_id
        visited: Set[str] = {current}

        for _ in range(max_hops):
            supersedes = self.get_supersedes(current)
            if not supersedes:
                break

            # Take the most confident superseded edge
            edges = self.get_outgoing(current, CounterfactualEdgeType.SUPERSEDES)
            best_edge = max(edges, key=lambda e: e.confidence)
            next_id = best_edge.target_episode_id

            if next_id in visited:
                logger.warning(f"Cycle detected in supersedes chain at {next_id}")
                break

            visited.add(next_id)
            current = next_id

        return current

    def find_chain(
        self,
        start_id: str,
        end_id: str,
        edge_types: Optional[Set[CounterfactualEdgeType]] = None,
        max_hops: int = 5,
    ) -> Optional[CounterfactualChain]:
        """
        Find a counterfactual chain between two episodes.

        Uses BFS to find the shortest path.

        Args:
            start_id: Starting episode ID
            end_id: Target episode ID
            edge_types: Optional filter for edge types to follow
            max_hops: Maximum chain length

        Returns:
            CounterfactualChain if found, None otherwise
        """
        if start_id == end_id:
            return CounterfactualChain(edges=[], start_episode_id=start_id, end_episode_id=end_id)

        # BFS
        queue: List[Tuple[str, List[CounterfactualEdge]]] = [(start_id, [])]
        visited: Set[str] = {start_id}

        while queue:
            current, path = queue.pop(0)

            if len(path) >= max_hops:
                continue

            for edge in self.get_outgoing(current):
                if edge_types and edge.edge_type not in edge_types:
                    continue

                target = edge.target_episode_id

                if target == end_id:
                    return CounterfactualChain(
                        edges=path + [edge],
                        start_episode_id=start_id,
                        end_episode_id=end_id,
                    )

                if target not in visited:
                    visited.add(target)
                    queue.append((target, path + [edge]))

        return None

    def get_all_edges(self) -> List[CounterfactualEdge]:
        """Get all edges in the graph."""
        all_edges = []
        for edges in self._outgoing.values():
            all_edges.extend(edges)
        return all_edges

    def to_dict(self) -> Dict[str, Any]:
        """Serialize the graph to a dictionary."""
        return {
            "edges": [e.to_dict() for e in self.get_all_edges()],
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "CounterfactualGraph":
        """Deserialize from a dictionary."""
        graph = cls()
        for edge_data in data.get("edges", []):
            edge = CounterfactualEdge.from_dict(edge_data)
            graph.add_edge(edge)
        return graph


class CounterfactualLinker:
    """
    Automatically creates counterfactual edges based on detected relationships.

    Integrates with ContradictionDetector to create edges when contradictions
    are found.
    """

    def __init__(
        self,
        graph: Optional[CounterfactualGraph] = None,
        auto_link_contradictions: bool = True,
        supersede_on_update: bool = True,
    ):
        """
        Initialize the linker.

        Args:
            graph: Counterfactual graph to add edges to
            auto_link_contradictions: Automatically link contradicting episodes
            supersede_on_update: Create supersedes edges for updates
        """
        self.graph = graph or CounterfactualGraph()
        self.auto_link_contradictions = auto_link_contradictions
        self.supersede_on_update = supersede_on_update

    def link_contradiction(
        self,
        episode_a_id: str,
        episode_b_id: str,
        confidence: float = 1.0,
        reason: Optional[str] = None,
    ) -> CounterfactualEdge:
        """Create a bidirectional contradicts edge."""
        edge = CounterfactualEdge(
            source_episode_id=episode_a_id,
            target_episode_id=episode_b_id,
            edge_type=CounterfactualEdgeType.CONTRADICTS,
            confidence=confidence,
            reason=reason,
        )
        self.graph.add_edge(edge)
        return edge

    def link_supersedes(
        self,
        new_episode_id: str,
        old_episode_id: str,
        confidence: float = 1.0,
        reason: Optional[str] = None,
    ) -> CounterfactualEdge:
        """Create a supersedes edge (new supersedes old)."""
        edge = CounterfactualEdge(
            source_episode_id=new_episode_id,
            target_episode_id=old_episode_id,
            edge_type=CounterfactualEdgeType.SUPERSEDES,
            confidence=confidence,
            reason=reason or "Updated information",
        )
        self.graph.add_edge(edge)
        return edge

    def link_obsoletes(
        self,
        new_episode_id: str,
        obsolete_episode_id: str,
        confidence: float = 1.0,
        reason: Optional[str] = None,
    ) -> CounterfactualEdge:
        """Create an obsoletes edge."""
        edge = CounterfactualEdge(
            source_episode_id=new_episode_id,
            target_episode_id=obsolete_episode_id,
            edge_type=CounterfactualEdgeType.OBSOLETES,
            confidence=confidence,
            reason=reason or "Information obsoleted",
        )
        self.graph.add_edge(edge)
        return edge

    def link_supports(
        self,
        supporting_id: str,
        supported_id: str,
        confidence: float = 1.0,
        reason: Optional[str] = None,
    ) -> CounterfactualEdge:
        """Create a supports edge."""
        edge = CounterfactualEdge(
            source_episode_id=supporting_id,
            target_episode_id=supported_id,
            edge_type=CounterfactualEdgeType.SUPPORTS,
            confidence=confidence,
            reason=reason,
        )
        self.graph.add_edge(edge)
        return edge

    def process_contradiction_result(
        self,
        result: Any,  # ContradictionResult
    ) -> List[CounterfactualEdge]:
        """
        Process contradiction detection results and create edges.

        Args:
            result: ContradictionResult from ContradictionDetector

        Returns:
            List of created edges
        """
        if not self.auto_link_contradictions:
            return []

        edges = []
        for contradiction in result.contradictions:
            edge = self.link_contradiction(
                episode_a_id=contradiction.episode_a_id,
                episode_b_id=contradiction.episode_b_id,
                confidence=contradiction.confidence,
                reason=contradiction.explanation,
            )
            edges.append(edge)

        return edges

# ------------------------------------------------------------------ #
# API Helpers (Exported for API Routes)
# ------------------------------------------------------------------ #


COUNTERFACTUAL_RELATIONS = {
    "contradicts",
    "supersedes",
    "obsoletes",
    "revises",
    "supports",
    "weakens",
}


def add_counterfactual_edge(
    graph_store: Any,
    source_episode_id: str,
    target_episode_id: str,
    relation: str,
    namespace: Optional[str] = None,
    weight: float = 1.0,
    metadata: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
    """
    Add a counterfactual edge to the graph store.

    Args:
        graph_store: The GraphStore instance
        source_episode_id: Source episode ID
        target_episode_id: Target episode ID
        relation: Relation type (must be one of COUNTERFACTUAL_RELATIONS)
        namespace: Optional namespace
        weight: Edge weight (confidence)
        metadata: Optional metadata

    Returns:
        The attributes dictionary of the created edge
    """
    if relation not in COUNTERFACTUAL_RELATIONS:
        raise ValueError(f"Invalid counterfactual relation: {relation}")

    attrs = {
        "relation": relation,
        "relation_type": relation,
        "weight": weight,
        "type": "counterfactual",
        "created_at": datetime.utcnow().isoformat(),
    }
    if metadata:
        attrs.update(metadata)

    # Use add_structural_edge if available, otherwise fallback might be needed
    # but GraphStore guarantees add_structural_edge exists (even if no-op)
    if hasattr(graph_store, "add_structural_edge"):
        graph_store.add_structural_edge(
            source_episode_id,
            target_episode_id,
            attrs,
            namespace=namespace
        )

    return attrs
