"""Graph route execution engine for reactive configurations.

Extracted from ``ReactiveCleverAgentsApp`` to keep ``application.py``
under the 500-line guideline (CONTRIBUTING.md).  The ``GraphExecutor``
class encapsulates iterative message-router graph traversal: topology
parsing, context initialisation, rule matching, and chained-edge
following.
"""

from __future__ import annotations

import logging
import re
from collections.abc import Callable
from typing import Any

from cleveragents.reactive.config_parser import ReactiveConfig
from cleveragents.reactive.route import RouteConfig
from cleveragents.reactive.stream_router import SimpleLLMAgent, SimpleToolAgent

logger = logging.getLogger(__name__)


class GraphExecutor:
    """Execute graph routes by following message_router rules iteratively.

    Args:
        agents: Registered agent instances (keyed by name).
        config: The parsed ``ReactiveConfig``, or ``None``.
    """

    def __init__(
        self,
        agents: dict[str, Any],
        config: ReactiveConfig | None = None,
    ) -> None:
        self._agents = agents
        self._config = config

    # ------------------------------------------------------------------
    # Public entry point
    # ------------------------------------------------------------------

    def execute(self, prompt: str, route: RouteConfig) -> str:
        """Execute a graph route and return the final output string."""
        router_node, router_rules, node_actor_map, node_edges = self._parse_topology(
            route
        )

        if not router_rules:
            return ""

        context = self._initialize_context()
        max_iterations = 50
        current_message = prompt

        def _condition_matches(condition: Any, ctx: dict[str, Any]) -> bool:
            if not condition:
                return True
            if isinstance(condition, dict) and condition.get("type") == "context_value":
                key = condition.get("key")
                if not key:
                    return False
                return ctx.get(key) == condition.get("value")
            return True

        def _select_targets(source_node: str) -> list[str]:
            edges = node_edges.get(source_node, [])
            return [
                e["target"]
                for e in edges
                if _condition_matches(e.get("condition"), context)
            ]

        workflow_hops = 0
        for _ in range(max_iterations):
            target_node, extracted_message, matched_rule = self._match_router_rule(
                current_message, router_rules
            )

            if target_node == "workflow_controller":
                is_default_rule = (
                    matched_rule
                    and matched_rule.get("match_type") in {"prefix", "suffix"}
                    and matched_rule.get("pattern", "") == ""
                )
                if workflow_hops > 0 and is_default_rule:
                    return extracted_message

            if target_node is None or target_node == "end":
                return extracted_message

            actor_name = node_actor_map.get(target_node, target_node)
            agent = self._agents.get(actor_name)

            if agent is None:
                logger.warning(
                    "Agent '%s' not found for node '%s'", actor_name, target_node
                )
                return extracted_message

            if target_node == "workflow_controller":
                workflow_hops += 1

            current_message = self._invoke_agent(agent, extracted_message, context)

            if not current_message:
                return ""

            next_targets = _select_targets(target_node)
            if router_node in next_targets:
                continue
            elif "end" in next_targets:
                return current_message
            elif next_targets:
                current_message, should_return = self._follow_chained_edges(
                    next_targets,
                    current_message,
                    context,
                    node_actor_map,
                    self._agents,
                    router_node,
                    _select_targets,
                )
                if should_return:
                    return current_message
                continue
            else:
                continue

        logger.warning("Graph route execution hit max iterations (%d)", max_iterations)
        return current_message

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _invoke_agent(agent: Any, message: str, context: dict[str, Any]) -> str:
        """Execute an agent and return the string result."""
        if isinstance(agent, (SimpleToolAgent | SimpleLLMAgent)):
            result = agent.process(message, context=context)
        elif hasattr(agent, "process"):
            result = agent.process(message)
        elif callable(agent):
            result = agent(message)
        else:
            result = message
        return str(result) if result else ""

    @staticmethod
    def _parse_topology(
        route: RouteConfig,
    ) -> tuple[
        str | None,
        list[dict[str, Any]],
        dict[str, str],
        dict[str, list[dict[str, Any]]],
    ]:
        """Parse route nodes/edges into router node, rules, actor map, and edge map."""
        router_node: str | None = None
        router_rules: list[dict[str, Any]] = []
        node_actor_map: dict[str, str] = {}
        node_edges: dict[str, list[dict[str, Any]]] = {}

        for node_name, node_data in route.nodes.items():
            node_type = node_data.get("type", "")
            if node_type == "message_router":
                router_node = node_name
                router_rules = node_data.get("rules", [])
            elif node_type == "actor":
                actor_name = node_data.get("actor", node_name)
                node_actor_map[node_name] = actor_name

        for edge_data in route.edges:
            if isinstance(edge_data, dict):
                source = edge_data.get("source", "")
                target = edge_data.get("target", "")
                condition = edge_data.get("condition")
            else:
                source = getattr(edge_data, "source", "")
                target = getattr(edge_data, "target", "")
                condition = getattr(edge_data, "condition", None)
            if source and target:
                node_edges.setdefault(source, []).append(
                    {"target": target, "condition": condition}
                )

        return router_node, router_rules, node_actor_map, node_edges

    def _initialize_context(self) -> dict[str, Any]:
        """Ensure global context has required default keys for graph execution."""
        context = self._config.global_context if self._config else {}
        default_stage_order = [
            "intro",
            "discovery",
            "brainstorming",
            "vetting",
            "structure",
            "section_writing",
            "paper_review",
            "latex_generation",
        ]
        stage_order = context.get("stage_order")
        requires_refresh = not isinstance(stage_order, list) or any(
            stage not in (stage_order or [])
            for stage in ("section_writing", "paper_review", "latex_generation")
        )
        if requires_refresh:
            context["stage_order"] = list(default_stage_order)
        if "writing_stage" not in context or context.get("writing_stage") not in (
            context.get("stage_order") or []
        ):
            context["writing_stage"] = "intro"
        if "paper_details" not in context:
            context["paper_details"] = {
                "topic": None,
                "length": None,
                "audience": None,
                "publication": None,
                "format": None,
                "other": None,
            }
        return context

    @staticmethod
    def _match_router_rule(
        message: str, router_rules: list[dict[str, Any]]
    ) -> tuple[str | None, str, dict[str, Any] | None]:
        """Match message against router rules.

        Returns (target_node, extracted_message, matched_rule).
        """
        for rule in router_rules:
            match_type = rule.get("match_type", "prefix")
            pattern = rule.get("pattern", "")
            rule_target = rule.get("target")
            extract = rule.get("extract_message", False)
            separator = rule.get("separator", ":")

            if not rule_target:
                continue

            matched = False
            if match_type == "prefix":
                matched = pattern == "" or message.startswith(pattern)
            elif match_type == "contains":
                matched = bool(pattern and pattern in message)
            elif match_type == "suffix":
                matched = pattern == "" or message.endswith(pattern)

            if matched:
                extracted = message
                if extract and separator and separator in message:
                    idx = message.index(separator)
                    extracted = message[idx + len(separator) :]
                return rule_target, extracted, rule

        return None, message, None

    @staticmethod
    def _follow_chained_edges(
        next_targets: list[str],
        current_message: str,
        context: dict[str, Any],
        node_actor_map: dict[str, str],
        agents: dict[str, Any],
        router_node: str | None,
        select_targets_fn: Callable[[str], list[str]],
    ) -> tuple[str, bool]:
        """Follow chained edges from a node.

        Returns (message, should_return_to_caller).
        """
        next_node = next_targets[0]
        while next_node:
            if next_node == "end":
                return current_message, True
            next_actor = node_actor_map.get(next_node, next_node)
            next_agent = agents.get(next_actor)
            if next_agent is not None:
                result = GraphExecutor._invoke_agent(
                    next_agent, current_message, context
                )
                current_message = result if result else current_message
            chained_targets = select_targets_fn(next_node)
            if router_node in chained_targets:
                return current_message, False
            if "end" in chained_targets:
                return current_message, True
            if not chained_targets:
                return current_message, False
            next_node = chained_targets[0]
        return current_message, False

    @staticmethod
    def strip_routing_prefixes(output: str) -> str:
        """Remove routing prefixes from output string.

        Strips known prefixes (``DISCOVERY_RESPONSE:``, ``SET_TOPIC:``,
        etc.) and generic ``ROUTE_*:`` patterns.
        """
        known_prefixes = [
            "DISCOVERY_RESPONSE:",
            "GOTO_BRAINSTORMING:",
            "SET_TOPIC:",
            "COMMAND_OUTPUT:",
        ]
        for prefix in known_prefixes:
            if output.startswith(prefix):
                output = output[len(prefix) :]
        route_match = re.match(r"^ROUTE_[A-Z_]+:", output)
        if route_match:
            output = output[route_match.end() :]
        return output

    @staticmethod
    def strip_routing_prefixes_multiline(output: str) -> str:
        """Remove routing prefixes from each line of a multi-line output."""
        known_prefixes = [
            "DISCOVERY_RESPONSE:",
            "GOTO_BRAINSTORMING:",
            "SET_TOPIC:",
            "COMMAND_OUTPUT:",
        ]
        cleaned_lines = []
        for line in output.split("\n"):
            for prefix in known_prefixes:
                if line.startswith(prefix):
                    line = line[len(prefix) :]
                    break
            else:
                route_match = re.match(r"^ROUTE_[A-Z_]+:", line)
                if route_match:
                    line = line[route_match.end() :]
            cleaned_lines.append(line)
        return "\n".join(cleaned_lines)
