"""
Transformation passes on IR.

Optimizations and validations that work on ManifestIR.
These are applied BETWEEN parsing and generation.

Each pass declares metadata including:
- name: Unique identifier
- description: What the pass does
- requires: Passes that must run before this one
- invalidates: Analyses invalidated by this pass
- preserves: Analyses preserved by this pass
"""

import logging
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import List, Optional, Set

from . import EdgeTrigger, GraphIR, ManifestIR, NodeKind

logger = logging.getLogger(__name__)


@dataclass
class PassMetadata:
    """Metadata for transformation pass."""

    name: str
    description: str
    requires: Set[str] = field(default_factory=set)  # Must run after these
    invalidates: Set[str] = field(default_factory=set)  # Invalidates these passes
    preserves: Set[str] = field(default_factory=set)  # Preserves these analyses


class Transform(ABC):
    """Base class for IR transformations with metadata."""

    @property
    @abstractmethod
    def name(self) -> str:
        """Transform name for logging."""
        ...

    @property
    def metadata(self) -> Optional[PassMetadata]:
        """Return pass metadata for dependency resolution."""
        return PassMetadata(
            name=self.name,
            description=f"Transform: {self.name}",
            requires=set(),
            invalidates=set(),
            preserves=set(),
        )

    @abstractmethod
    def apply(self, ir: ManifestIR) -> ManifestIR:
        """Apply transformation to IR."""
        ...


class DeadNodeElimination(Transform):
    """Remove unreachable nodes from graphs."""

    @property
    def name(self) -> str:
        return "dead-node-elimination"

    @property
    def metadata(self) -> PassMetadata:
        return PassMetadata(
            name=self.name,
            description="Remove unreachable nodes from entry point",
            requires=set(),  # Can run first
            invalidates={"edge-analysis"},  # May change edges
            preserves={"cycle-free"},  # Preserves cycle-freedom
        )

    def apply(self, ir: ManifestIR) -> ManifestIR:
        """Remove unreachable nodes."""
        for graph in ir.graphs:
            unreachable = graph.find_unreachable_nodes()

            if unreachable:
                logger.info(
                    f"[{self.name}] Removing {len(unreachable)} unreachable nodes: {unreachable}"
                )
                graph.nodes = [n for n in graph.nodes if n.id not in unreachable]
                graph.edges = [
                    e
                    for e in graph.edges
                    if e.from_node not in unreachable and e.to_node not in unreachable
                ]
                graph._build_indexes()

        return ir


class EdgeDeduplication(Transform):
    """Remove duplicate edges."""

    @property
    def name(self) -> str:
        return "edge-deduplication"

    @property
    def metadata(self) -> PassMetadata:
        return PassMetadata(
            name=self.name,
            description="Remove duplicate edges between nodes",
            requires={"dead-node-elimination"},  # Run after dead code removal
            invalidates=set(),
            preserves={"dag", "cycle-free"},
        )

    def apply(self, ir: ManifestIR) -> ManifestIR:
        """Deduplicate edges."""
        for graph in ir.graphs:
            seen: Set[tuple] = set()
            unique_edges = []

            for edge in graph.edges:
                key = (
                    edge.from_node,
                    edge.to_node,
                    edge.condition.trigger,
                    edge.condition.route,
                )
                if key not in seen:
                    seen.add(key)
                    unique_edges.append(edge)

            removed = len(graph.edges) - len(unique_edges)
            if removed > 0:
                logger.info(f"[{self.name}] Removed {removed} duplicate edges")
                graph.edges = unique_edges
                graph._build_indexes()

        return ir


class ConditionSimplification(Transform):
    """Simplify edge conditions."""

    @property
    def name(self) -> str:
        return "condition-simplification"

    @property
    def metadata(self) -> PassMetadata:
        return PassMetadata(
            name=self.name,
            description="Simplify boolean expressions in edge conditions",
            requires={"edge-deduplication"},  # Run after deduplication
            invalidates=set(),
            preserves={"dag", "cycle-free"},
        )

    def apply(self, ir: ManifestIR) -> ManifestIR:
        """Simplify conditions."""
        simplified_count = 0

        for graph in ir.graphs:
            for edge in graph.edges:
                if edge.condition.expression:
                    original = edge.condition.expression
                    simplified = self._simplify(original)
                    if simplified != original:
                        edge.condition.expression = simplified
                        simplified_count += 1

        if simplified_count > 0:
            logger.info(
                f"[{self.name}] Simplified {simplified_count} condition expressions"
            )

        return ir

    def _simplify(self, expr: str) -> str:
        """Simplify boolean expression."""
        # Basic simplifications
        simplifications = [
            ("True and ", ""),
            (" and True", ""),
            ("False or ", ""),
            (" or False", ""),
            ("not not ", ""),
        ]

        result = expr
        for pattern, replacement in simplifications:
            result = result.replace(pattern, replacement)

        return result.strip()


class RouterValidation(Transform):
    """Validate router nodes have corresponding router definitions."""

    @property
    def name(self) -> str:
        return "router-validation"

    @property
    def metadata(self) -> PassMetadata:
        return PassMetadata(
            name=self.name,
            description="Validate router references exist",
            requires={"dead-node-elimination"},  # Only validate reachable nodes
            invalidates=set(),
            preserves={"dag", "cycle-free", "edge-analysis"},
        )

    def apply(self, ir: ManifestIR) -> ManifestIR:
        """Validate router references."""
        router_names = {r.name for r in ir.routers}

        for graph in ir.graphs:
            for node in graph.nodes:
                if node.kind == NodeKind.ROUTER and node.router_ref:
                    if node.router_ref not in router_names:
                        logger.warning(
                            f"[{self.name}] Node '{node.id}' references undefined router: {node.router_ref}"
                        )

        return ir


class ToolValidation(Transform):
    """Validate tool nodes have corresponding tool definitions."""

    @property
    def name(self) -> str:
        return "tool-validation"

    @property
    def metadata(self) -> PassMetadata:
        return PassMetadata(
            name=self.name,
            description="Validate tool references exist",
            requires={"dead-node-elimination"},  # Only validate reachable nodes
            invalidates=set(),
            preserves={"dag", "cycle-free", "edge-analysis"},
        )

    def apply(self, ir: ManifestIR) -> ManifestIR:
        """Validate tool references."""
        tool_names = {t.name for t in ir.tools}

        for graph in ir.graphs:
            for node in graph.nodes:
                if node.kind == NodeKind.TOOL and node.tool_ref:
                    if node.tool_ref not in tool_names:
                        logger.warning(
                            f"[{self.name}] Node '{node.id}' references undefined tool: {node.tool_ref}"
                        )

        return ir


class CycleDetection(Transform):
    """Detect and warn about cycles in graph."""

    @property
    def name(self) -> str:
        return "cycle-detection"

    @property
    def metadata(self) -> PassMetadata:
        return PassMetadata(
            name=self.name,
            description="Detect cycles in graph (warning only)",
            requires={"dead-node-elimination", "edge-deduplication"},  # Run late
            invalidates=set(),
            preserves={"dag", "edge-analysis"},
        )

    def apply(self, ir: ManifestIR) -> ManifestIR:
        """Detect cycles."""
        for graph in ir.graphs:
            cycles = self._find_cycles(graph)
            if cycles:
                logger.warning(
                    f"[{self.name}] Graph '{graph.name}' contains {len(cycles)} cycle(s): {cycles}"
                )

        return ir

    def _find_cycles(self, graph: GraphIR) -> List[List[str]]:
        """Find all cycles in graph using DFS."""
        cycles = []
        visited = set()
        rec_stack = set()
        path = []

        def dfs(node_id: str) -> bool:
            visited.add(node_id)
            rec_stack.add(node_id)
            path.append(node_id)

            for edge in graph.get_outgoing_edges(node_id):
                next_node = edge.to_node
                if next_node not in visited:
                    if dfs(next_node):
                        return True
                elif next_node in rec_stack:
                    # Found cycle
                    cycle_start = path.index(next_node)
                    cycle = path[cycle_start:] + [next_node]
                    cycles.append(cycle)

            path.pop()
            rec_stack.remove(node_id)
            return False

        for node in graph.nodes:
            if node.id not in visited:
                dfs(node.id)

        return cycles


class EmptyGraphRemoval(Transform):
    """Remove empty graphs from manifest."""

    @property
    def name(self) -> str:
        return "empty-graph-removal"

    @property
    def metadata(self) -> PassMetadata:
        return PassMetadata(
            name=self.name,
            description="Remove graphs with no nodes",
            requires={"dead-node-elimination"},  # May create empty graphs
            invalidates=set(),
            preserves={"dag", "cycle-free"},
        )

    def apply(self, ir: ManifestIR) -> ManifestIR:
        """Remove empty graphs."""
        non_empty = [g for g in ir.graphs if g.nodes]
        removed = len(ir.graphs) - len(non_empty)

        if removed > 0:
            logger.info(f"[{self.name}] Removed {removed} empty graph(s)")
            ir.graphs = non_empty
            ir._build_indexes()

        return ir


# ===== ADVANCED OPTIMIZATION PASSES =====


class ConstantFolding(Transform):
    """
    Evaluate constant expressions at compile time.

    For agent graphs, this means:
    - Evaluate static conditions (True/False in expressions)
    - Remove edges that can never be taken
    - Simplify router conditions with known values
    """

    @property
    def name(self) -> str:
        return "constant-folding"

    @property
    def metadata(self) -> PassMetadata:
        return PassMetadata(
            name=self.name,
            description="Evaluate constant expressions at compile time",
            requires={"condition-simplification"},
            invalidates={"edge-analysis"},
            preserves={"cycle-free"},
        )

    def apply(self, ir: ManifestIR) -> ManifestIR:
        """Fold constant expressions."""
        folded_count = 0

        for graph in ir.graphs:
            edges_to_remove = []

            for edge in graph.edges:
                if edge.condition.expression:
                    result = self._try_evaluate(edge.condition.expression)

                    if result is True:
                        # Always taken - simplify to no condition
                        edge.condition.expression = None
                        folded_count += 1
                    elif result is False:
                        # Never taken - mark for removal
                        edges_to_remove.append(edge)
                        folded_count += 1

            # Remove never-taken edges
            if edges_to_remove:
                graph.edges = [e for e in graph.edges if e not in edges_to_remove]
                graph._build_indexes()
                logger.info(
                    f"[{self.name}] Removed {len(edges_to_remove)} dead edges"
                )

        if folded_count > 0:
            logger.info(f"[{self.name}] Folded {folded_count} constant expressions")

        return ir

    def _try_evaluate(self, expr: str) -> Optional[bool]:
        """Try to evaluate expression as constant. Returns None if not constant."""
        # Simple constant patterns
        expr_stripped = expr.strip()

        if expr_stripped == "True":
            return True
        elif expr_stripped == "False":
            return False

        # Try safe eval for simple expressions
        try:
            result = eval(expr_stripped, {"__builtins__": {}}, {})
            if isinstance(result, bool):
                return result
        except Exception:
            pass

        return None


class CommonSubexpressionElimination(Transform):
    """
    Eliminate duplicate tool calls and expressions.

    For agent graphs, this means:
    - Detect identical tool nodes with same config
    - Merge edges to single tool instance
    - Cache repeated computations
    """

    @property
    def name(self) -> str:
        return "cse"

    @property
    def metadata(self) -> PassMetadata:
        return PassMetadata(
            name=self.name,
            description="Eliminate duplicate tool calls",
            requires={"dead-node-elimination", "edge-deduplication"},
            invalidates={"edge-analysis"},
            preserves={"cycle-free"},
        )

    def apply(self, ir: ManifestIR) -> ManifestIR:
        """Eliminate common subexpressions."""
        merged_count = 0

        for graph in ir.graphs:
            # Find duplicate tool nodes
            tool_signatures: dict = {}

            for node in graph.nodes:
                if node.kind == NodeKind.TOOL and node.tool_ref:
                    sig = (node.tool_ref, str(sorted(node.config.items())))

                    if sig in tool_signatures:
                        # Found duplicate - redirect edges
                        original = tool_signatures[sig]
                        self._redirect_edges(graph, node.id, original.id)
                        merged_count += 1
                    else:
                        tool_signatures[sig] = node

            # Remove merged nodes
            if merged_count > 0:
                remaining_ids = {sig[0] for sig in tool_signatures.values()}
                # Actually we need to track which nodes we kept
                # This is simplified - in production you'd track redirects

        if merged_count > 0:
            logger.info(f"[{self.name}] Merged {merged_count} duplicate tool calls")

        return ir

    def _redirect_edges(self, graph: GraphIR, from_id: str, to_id: str) -> None:
        """Redirect all edges from one node to another."""
        for edge in graph.edges:
            if edge.to_node == from_id:
                edge.to_node = to_id
            if edge.from_node == from_id:
                edge.from_node = to_id


class InlineSmallGraphs(Transform):
    """
    Inline single-node subgraphs into callers.

    For agent graphs with multiple graphs (subgraphs), this:
    - Identifies graphs with only 1-2 nodes
    - Inlines them at call sites
    - Reduces graph traversal overhead
    """

    @property
    def name(self) -> str:
        return "inline-small-graphs"

    @property
    def metadata(self) -> PassMetadata:
        return PassMetadata(
            name=self.name,
            description="Inline single-node subgraphs",
            requires={"dead-node-elimination"},
            invalidates={"graph-analysis"},
            preserves={"cycle-free"},
        )

    def __init__(self, max_nodes: int = 2):
        """
        Initialize with max node threshold.

        Args:
            max_nodes: Maximum nodes in graph to consider for inlining
        """
        self.max_nodes = max_nodes

    def apply(self, ir: ManifestIR) -> ManifestIR:
        """Inline small graphs."""
        if len(ir.graphs) <= 1:
            return ir  # Nothing to inline

        inlined_count = 0
        graphs_to_remove = []

        # Find small graphs
        small_graphs = [g for g in ir.graphs if len(g.nodes) <= self.max_nodes]

        # For now, just log candidates (full inlining requires call graph analysis)
        for graph in small_graphs:
            if graph.name != "main":  # Don't inline main
                logger.debug(
                    f"[{self.name}] Candidate for inlining: {graph.name} "
                    f"({len(graph.nodes)} nodes)"
                )

        if inlined_count > 0:
            logger.info(f"[{self.name}] Inlined {inlined_count} small graphs")

        return ir


# ===== TRANSFORM PIPELINE =====


class TransformPipeline:
    """
    Pipeline of IR transformations.

    Transforms are applied in order.
    """

    def __init__(self, transforms: List[Transform] = None):
        self.transforms = transforms or []

    def add(self, transform: Transform) -> "TransformPipeline":
        """Add transform to pipeline."""
        self.transforms.append(transform)
        return self

    def apply(self, ir: ManifestIR) -> ManifestIR:
        """Apply all transforms."""
        for transform in self.transforms:
            logger.debug(f"Applying transform: {transform.name}")
            ir = transform.apply(ir)
        return ir


# ===== DEFAULT PIPELINES =====


def create_optimization_pipeline() -> TransformPipeline:
    """Create default optimization pipeline."""
    return TransformPipeline(
        [
            DeadNodeElimination(),
            EdgeDeduplication(),
            ConditionSimplification(),
            EmptyGraphRemoval(),
        ]
    )


def create_validation_pipeline() -> TransformPipeline:
    """Create validation pipeline."""
    return TransformPipeline(
        [
            RouterValidation(),
            ToolValidation(),
            CycleDetection(),
        ]
    )


def create_full_pipeline() -> TransformPipeline:
    """Create full pipeline (validation + optimization)."""
    return TransformPipeline(
        [
            # Validations first
            RouterValidation(),
            ToolValidation(),
            CycleDetection(),
            # Then optimizations
            DeadNodeElimination(),
            EdgeDeduplication(),
            ConditionSimplification(),
            EmptyGraphRemoval(),
        ]
    )


def optimize(ir: ManifestIR) -> ManifestIR:
    """
    Apply default optimization passes.

    Convenience function for simple usage.
    """
    pipeline = create_optimization_pipeline()
    return pipeline.apply(ir)


def validate(ir: ManifestIR) -> ManifestIR:
    """
    Apply validation passes.

    Convenience function for simple usage.
    """
    pipeline = create_validation_pipeline()
    return pipeline.apply(ir)

