"""
Belief Revision System - Inference Engine

Graph traversal and path-finding for belief networks.
Supports ancestry tracing, reachability queries, and path explanation.
"""
from __future__ import annotations

import heapq
from dataclasses import dataclass
from typing import Any, Callable, Dict, List, Optional, Set, Tuple

from .storage import CASStore


# =============================================================================
# Cost Functions (Pluggable)
# =============================================================================

# Default relation penalties (lower = stronger relationship)
DEFAULT_RELATION_PENALTIES = {
    "direct_descent": 0.0,
    "derived_via": 0.5,
    "influence": 1.0,
    "stimulus_diffusion": 1.5,
}


def default_edge_cost(
    relation: str,
    tier: int,
    confidence: float,
    relation_penalties: Optional[Dict[str, float]] = None
) -> float:
    """
    Compute traversal cost for an edge.

    Lower cost = stronger, more direct relationship.

    Args:
        relation: Edge relation type
        tier: Relationship strength tier (0=strongest, 5=weakest)
        confidence: Confidence score (0-1)
        relation_penalties: Optional custom penalty mapping

    Returns:
        Edge traversal cost
    """
    penalties = relation_penalties or DEFAULT_RELATION_PENALTIES
    rel_pen = penalties.get(relation, 1.0)
    tier_pen = tier * 0.3
    conf_pen = (1.0 - confidence) * 1.5
    return rel_pen + tier_pen + conf_pen


# Type alias for edge cost functions
EdgeCostFn = Callable[[str, int, float], float]


# =============================================================================
# Path Results
# =============================================================================

@dataclass
class PathStep:
    """A single step in a traversal path."""
    parent: str
    child: str
    edge_id: str
    relation: str
    tier: int
    confidence: float


@dataclass
class AncestryResult:
    """Result of an ancestry path query."""
    target: str          # Starting node
    root: str            # Found root node
    steps: List[PathStep]  # Path from target to root
    score: float         # Total path cost (lower = better)
    explanation: str     # Human-readable explanation


# =============================================================================
# Path-Finding Algorithms
# =============================================================================

def best_path_to_any_root(
    store: CASStore,
    node_id: str,
    roots: List[str],
    max_hops: int = 20,
    edge_cost_fn: EdgeCostFn = default_edge_cost
) -> Optional[AncestryResult]:
    """
    Find the best (lowest cost) path from a node to any root.

    Uses Dijkstra's algorithm traversing edges upward (child → parent).

    Args:
        store: Storage backend
        node_id: Starting node
        roots: Candidate root nodes
        max_hops: Maximum traversal depth
        edge_cost_fn: Function to compute edge costs

    Returns:
        AncestryResult if path found, None otherwise
    """
    roots_set = set(roots)

    # Priority queue: (cost, hops, current_node)
    pq: List[Tuple[float, int, str]] = [(0.0, 0, node_id)]
    best_cost: Dict[str, float] = {node_id: 0.0}

    # Track path: current_node → (next_node_toward_root, edge_dict)
    prev: Dict[str, Tuple[str, Dict[str, Any]]] = {}

    while pq:
        cost, hops, cur = heapq.heappop(pq)

        # Skip if we've found a better path
        if cost != best_cost.get(cur, float("inf")):
            continue

        # Found a root!
        if cur in roots_set:
            steps = _reconstruct_path(node_id, cur, prev)
            explanation = _explain_path(store, node_id, cur, steps, cost)
            return AncestryResult(
                target=node_id,
                root=cur,
                steps=steps,
                score=cost,
                explanation=explanation
            )

        # Depth limit
        if hops >= max_hops:
            continue

        # Traverse incoming edges (parents of current node)
        for edge in store.list_edges_into(cur):
            parent = edge["parent_id"]
            c = edge_cost_fn(edge["relation"], edge["tier"], edge["confidence"])
            new_cost = cost + c

            if new_cost < best_cost.get(parent, float("inf")):
                best_cost[parent] = new_cost
                prev[cur] = (parent, edge)
                heapq.heappush(pq, (new_cost, hops + 1, parent))

    return None


def find_all_paths(
    store: CASStore,
    start: str,
    end: str,
    max_hops: int = 10,
    max_paths: int = 5
) -> List[List[PathStep]]:
    """
    Find all paths between two nodes (up to limits).

    Uses BFS with path tracking. Useful for understanding
    multiple derivation routes.

    Args:
        store: Storage backend
        start: Starting node
        end: Target node
        max_hops: Maximum path length
        max_paths: Maximum paths to return

    Returns:
        List of paths (each path is a list of PathSteps)
    """
    paths: List[List[PathStep]] = []

    # Queue: (current_node, path_so_far, visited)
    queue: List[Tuple[str, List[PathStep], Set[str]]] = [(start, [], {start})]

    while queue and len(paths) < max_paths:
        cur, path, visited = queue.pop(0)

        if cur == end:
            paths.append(path)
            continue

        if len(path) >= max_hops:
            continue

        # Try edges in both directions
        for edge in store.list_edges_from(cur):
            child = edge["child_id"]
            if child not in visited:
                step = PathStep(
                    parent=cur, child=child,
                    edge_id=edge["edge_id"],
                    relation=edge["relation"],
                    tier=edge["tier"],
                    confidence=edge["confidence"]
                )
                queue.append((child, path + [step], visited | {child}))

        for edge in store.list_edges_into(cur):
            parent = edge["parent_id"]
            if parent not in visited:
                step = PathStep(
                    parent=parent, child=cur,
                    edge_id=edge["edge_id"],
                    relation=edge["relation"],
                    tier=edge["tier"],
                    confidence=edge["confidence"]
                )
                queue.append((parent, path + [step], visited | {parent}))

    return paths


def is_reachable(
    store: CASStore,
    start: str,
    end: str,
    max_hops: int = 20,
    direction: str = "up"  # "up" (to parents), "down" (to children), "any"
) -> bool:
    """
    Check if end is reachable from start.

    Args:
        store: Storage backend
        start: Starting node
        end: Target node
        max_hops: Maximum traversal depth
        direction: Traversal direction

    Returns:
        True if path exists
    """
    visited: Set[str] = set()
    queue = [(start, 0)]

    while queue:
        cur, depth = queue.pop(0)

        if cur == end:
            return True

        if cur in visited or depth >= max_hops:
            continue
        visited.add(cur)

        if direction in ("up", "any"):
            for edge in store.list_edges_into(cur):
                queue.append((edge["parent_id"], depth + 1))

        if direction in ("down", "any"):
            for edge in store.list_edges_from(cur):
                queue.append((edge["child_id"], depth + 1))

    return False


def get_ancestors(
    store: CASStore,
    node_id: str,
    max_depth: int = 10
) -> Set[str]:
    """Get all ancestors of a node up to max_depth."""
    ancestors: Set[str] = set()
    queue = [(node_id, 0)]
    visited: Set[str] = set()

    while queue:
        cur, depth = queue.pop(0)

        if cur in visited or depth > max_depth:
            continue
        visited.add(cur)

        for edge in store.list_edges_into(cur):
            parent = edge["parent_id"]
            ancestors.add(parent)
            queue.append((parent, depth + 1))

    return ancestors


def get_descendants(
    store: CASStore,
    node_id: str,
    max_depth: int = 10
) -> Set[str]:
    """Get all descendants of a node up to max_depth."""
    descendants: Set[str] = set()
    queue = [(node_id, 0)]
    visited: Set[str] = set()

    while queue:
        cur, depth = queue.pop(0)

        if cur in visited or depth > max_depth:
            continue
        visited.add(cur)

        for edge in store.list_edges_from(cur):
            child = edge["child_id"]
            descendants.add(child)
            queue.append((child, depth + 1))

    return descendants


# =============================================================================
# Helper Functions
# =============================================================================

def _reconstruct_path(
    start: str,
    end: str,
    prev: Dict[str, Tuple[str, Dict[str, Any]]]
) -> List[PathStep]:
    """Reconstruct path from prev mapping."""
    steps: List[PathStep] = []
    cur = start

    while cur != end:
        if cur not in prev:
            break
        parent, edge = prev[cur]
        steps.append(PathStep(
            parent=parent,
            child=cur,
            edge_id=edge["edge_id"],
            relation=edge["relation"],
            tier=edge["tier"],
            confidence=edge["confidence"]
        ))
        cur = parent

    return list(reversed(steps))


def _explain_path(
    store: CASStore,
    target: str,
    root: str,
    steps: List[PathStep],
    score: float
) -> str:
    """Generate human-readable path explanation."""
    def get_name(node_id: str) -> str:
        try:
            return store.get_node_by_id(node_id)["name"]
        except Exception:
            return node_id

    lines = [
        f"Ancestry path for {get_name(target)} ({target})",
        f"→ root: {get_name(root)} ({root})",
        f"Score (lower is better): {score:.3f}",
        ""
    ]

    if not steps:
        lines.append("Already at root.")
        return "\n".join(lines)

    lines.append("Path (toward root):")
    for step in steps:
        lines.append(
            f"  {get_name(step.child)} ← {get_name(step.parent)} "
            f"[{step.relation}, tier={step.tier}, conf={step.confidence:.2f}]"
        )

    return "\n".join(lines)
