"""
TopologyEngine — unified orchestration entry point.

Provides a single interface over the three orchestration modes:
  - GRAPH: Static graph execution (GraphOrchestrator)
  - ROUTER: Dynamic routing (RouterOrchestrator)
  - ADAPTIVE: Self-organizing topology (AdaptiveOrchestrator)
"""

from __future__ import annotations

import logging
from enum import StrEnum
from typing import Any

from nexus.core.agent import Agent
from nexus.core.state import SharedState
from nexus.core.task import Task, TaskGraph, TaskResult
from nexus.orchestration.adaptive import AdaptiveOrchestrator
from nexus.orchestration.graph import GraphOrchestrator
from nexus.orchestration.router import RouterOrchestrator
from nexus.protocol.discovery import DiscoveryService

logger = logging.getLogger("nexus.orchestration.topology")


class TopologyMode(StrEnum):
    """The three orchestration modes."""

    GRAPH = "graph"
    ROUTER = "router"
    ADAPTIVE = "adaptive"


class TopologyEngine:
    """
    Unified orchestration engine — the single entry point for running
    multi-agent workflows in any of the three topology modes.

    Usage:
        engine = TopologyEngine(mode=TopologyMode.ADAPTIVE)
        engine.add_agent(researcher)
        engine.add_agent(writer)

        # Adaptive mode
        results = await engine.run_tasks(task_graph)

        # Router mode
        engine.set_mode(TopologyMode.ROUTER)
        engine.router.add_route("research", agent_id=researcher.id)
        result = await engine.run_single(task)

        # Graph mode
        engine.set_mode(TopologyMode.GRAPH)
        engine.graph.add_node("plan", planner_func)
        engine.graph.add_edge(START, "plan")
        state = await engine.run_graph(shared_state)
    """

    def __init__(
        self,
        mode: TopologyMode = TopologyMode.ADAPTIVE,
        discovery: DiscoveryService | None = None,
    ) -> None:
        self.mode = mode
        self.discovery = discovery or DiscoveryService()

        # Initialize sub-orchestrators
        self._graph: GraphOrchestrator[Any] | None = None
        self._router = RouterOrchestrator(discovery=self.discovery)
        self._adaptive = AdaptiveOrchestrator(discovery=self.discovery)

    def set_mode(self, mode: TopologyMode) -> None:
        """Switch the orchestration mode."""
        self.mode = mode
        logger.info(f"Topology mode set to: {mode.value}")

    @property
    def graph(self) -> GraphOrchestrator[Any]:
        """Access the graph orchestrator (creates one if needed)."""
        if self._graph is None:
            self._graph = GraphOrchestrator()
        return self._graph

    @property
    def router(self) -> RouterOrchestrator:
        return self._router

    @property
    def adaptive(self) -> AdaptiveOrchestrator:
        return self._adaptive

    def add_agent(self, agent: Agent) -> None:
        """Register an agent with all orchestrators."""
        self.discovery.register(agent)
        self._router.add_agent(agent.name, agent)
        self._adaptive.add_agent(agent)

    # ---- Execution methods ----

    async def run_single(self, task: Task) -> TaskResult:
        """
        Execute a single task using the current topology mode.
        """
        if self.mode == TopologyMode.ROUTER:
            return await self._router.route_and_execute(task)
        elif self.mode == TopologyMode.ADAPTIVE:
            from nexus.core.task import TaskGraph
            graph = TaskGraph()
            graph.add_task(task)
            results = await self._adaptive.run(graph)
            if task.id in results:
                return results[task.id]
            raise RuntimeError(f"Task '{task.name}' produced no result")
        else:
            raise ValueError(
                "run_single is not supported in GRAPH mode. "
                "Use run_graph() instead."
            )

    async def run_tasks(self, graph: TaskGraph) -> dict[str, TaskResult]:
        """
        Execute a task graph using the current topology mode.
        """
        if self.mode == TopologyMode.ADAPTIVE:
            return await self._adaptive.run(graph)
        elif self.mode == TopologyMode.ROUTER:
            results = {}
            for task in graph.topological_sort():
                # Inject dependency results
                for dep_id in task.depends_on:
                    dep_result = results.get(dep_id)
                    if dep_result:
                        task.input_data[f"dep_{dep_id}"] = dep_result.content
                result = await self._router.route_and_execute(task)
                results[task.id] = result
            return results
        else:
            raise ValueError(
                "run_tasks is not supported in GRAPH mode. "
                "Use run_graph() instead."
            )

    async def run_graph(self, state: SharedState[Any], max_steps: int = 100) -> SharedState[Any]:
        """
        Execute a static graph workflow.

        Only available in GRAPH mode.
        """
        if self.mode != TopologyMode.GRAPH:
            logger.warning(
                f"run_graph called in {self.mode.value} mode, "
                f"switching to GRAPH mode"
            )
            self.mode = TopologyMode.GRAPH

        return await self.graph.run(state, max_steps=max_steps)
