"""
ExecutionEngine — the core runtime that drives task execution through agents.

Design notes:
  - Orchestrates the full lifecycle: task → agent assignment → execution → checkpoint → result
  - Integrates with the PolicyChain for error recovery
  - Integrates with TokenBudgetManager for cost control
  - Integrates with CheckpointStore for durability
  - Handles handoffs (delegate, transfer, broadcast)
  - Supports parallel execution of independent tasks
"""

from __future__ import annotations

import asyncio
import logging
from typing import Any

from nexus.core.agent import Agent
from nexus.core.handoff import Handoff, HandoffMode
from nexus.core.message import Message
from nexus.core.task import Task, TaskGraph, TaskResult, TaskStatus
from nexus.protocol.bus import MessageBus
from nexus.protocol.discovery import DiscoveryService
from nexus.runtime.budget import BudgetConfig, TokenBudgetManager
from nexus.runtime.checkpoint import Checkpoint, CheckpointStore, InMemoryCheckpointStore
from nexus.runtime.policies import PolicyChain
from nexus.types import QualityRecord

logger = logging.getLogger("nexus.runtime.engine")


class ExecutionEngine:
    """
    The heart of the Nexus runtime — drives tasks through agents with
    full error recovery, budget control, checkpointing, and handoff support.

    Usage:
        engine = ExecutionEngine()

        # Register agents
        engine.register_agent(researcher)
        engine.register_agent(writer)

        # Execute a single task
        result = await engine.run_task(task, agent_id=researcher.id)

        # Execute a task graph (with dependencies and parallelism)
        results = await engine.run_graph(task_graph)
    """

    def __init__(
        self,
        checkpoint_store: CheckpointStore | None = None,
        budget_config: BudgetConfig | None = None,
        policy_chain: PolicyChain | None = None,
        message_bus: MessageBus | None = None,
        discovery: DiscoveryService | None = None,
    ) -> None:
        self.checkpoint_store = checkpoint_store or InMemoryCheckpointStore()
        self.budget = TokenBudgetManager(budget_config)
        self.policies = policy_chain or PolicyChain()
        self.bus = message_bus or MessageBus()
        self.discovery = discovery or DiscoveryService()

        self._running_tasks: dict[str, asyncio.Task[Any]] = {}

    def register_agent(self, agent: Agent) -> None:
        """Register an agent with the engine's discovery service."""
        self.discovery.register(agent)

    # ---- Single task execution ----

    async def run_task(
        self,
        task: Task,
        agent_id: str | None = None,
        checkpoint: bool = True,
    ) -> TaskResult:
        """
        Execute a single task with an agent.

        If agent_id is not specified, the engine uses the discovery service
        to find the best agent for the task.
        """
        # Resolve agent
        agent = self._resolve_agent(task, agent_id)
        if not agent:
            task.fail("No suitable agent found")
            raise RuntimeError(f"No agent available for task '{task.name}'")

        logger.info(f"Executing task '{task.name}' with agent '{agent.name}'")

        # Check budget
        if not self.budget.can_spend(estimated_tokens=1000):
            task.fail("Token budget exhausted")
            raise RuntimeError("Token budget exhausted")

        try:
            # Execute with policy chain
            result = await self.policies.execute(
                task=task,
                task_id=task.id,
                agent_id=agent.id,
                primary_func=agent.run,
                fallback_func=self._make_fallback_executor(task),
            )

            # Checkpoint on success
            if checkpoint:
                await self._save_checkpoint(task, agent.id)

            # Record budget usage
            await self.budget.record(
                task.token_usage,
                agent_id=agent.id,
                task_id=task.id,
            )

            # Record quality for the feedback loop
            self._record_quality(task, agent.id, result)

            return result

        except Exception:
            # Record failure
            self.discovery.record_failure(agent.id, task.id)
            if checkpoint:
                await self._save_checkpoint(task, agent.id)
            raise

    # ---- Task graph execution ----

    async def run_graph(
        self,
        graph: TaskGraph,
        checkpoint: bool = True,
    ) -> dict[str, TaskResult]:
        """
        Execute a task graph with parallel execution of independent tasks.

        Tasks are scheduled based on their dependency edges.
        Tasks with no unfinished dependencies run concurrently.
        """
        results: dict[str, TaskResult] = {}

        logger.info(
            f"Executing task graph with {len(graph.tasks)} tasks"
        )

        while not graph.is_complete:
            ready_tasks = graph.get_ready_tasks()
            if not ready_tasks:
                # Check for deadlock (all remaining tasks have unmet dependencies)
                remaining = [
                    t for t in graph.tasks.values() if not t.status.is_terminal
                ]
                if remaining:
                    failed = [t.id for t in remaining]
                    raise RuntimeError(
                        f"Task graph deadlock: tasks {failed} have unmet dependencies"
                    )
                break

            # Execute ready tasks in parallel
            coros = []
            for task in ready_tasks:
                coros.append(self._run_graph_task(task, graph, results, checkpoint))

            await asyncio.gather(*coros, return_exceptions=True)

        if graph.has_failures:
            failed_tasks = [
                t.name for t in graph.tasks.values()
                if t.status == TaskStatus.FAILED
            ]
            logger.warning(f"Task graph completed with failures: {failed_tasks}")

        return results

    async def _run_graph_task(
        self,
        task: Task,
        graph: TaskGraph,
        results: dict[str, TaskResult],
        checkpoint: bool,
    ) -> None:
        """Execute a single task within a graph context."""
        # Inject dependency results into the task's input
        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

        try:
            result = await self.run_task(task, checkpoint=checkpoint)
            results[task.id] = result
        except Exception as exc:
            logger.error(f"Task '{task.name}' failed: {exc}")
            task.fail(str(exc))

    # ---- Handoff execution ----

    async def execute_handoff(self, handoff: Handoff, source_task: Task) -> Any:
        """
        Execute a handoff between agents.

        Handles all three modes: delegate, transfer, broadcast.
        """
        if handoff.mode == HandoffMode.DELEGATE:
            return await self._delegate(handoff, source_task)
        elif handoff.mode == HandoffMode.TRANSFER:
            return await self._transfer(handoff, source_task)
        elif handoff.mode == HandoffMode.BROADCAST:
            return await self._broadcast(handoff, source_task)
        else:
            raise ValueError(f"Unknown handoff mode: {handoff.mode}")

    async def _delegate(self, handoff: Handoff, source_task: Task) -> TaskResult:
        """Delegate: send a sub-task and wait for the result."""
        sub_task = Task(
            name=f"delegated_from_{source_task.name}",
            description=handoff.message,
            input_data=handoff.context_data,
        )
        result = await self.run_task(sub_task, agent_id=handoff.target_agent_id)
        handoff.complete(result.content)
        return result

    async def _transfer(self, handoff: Handoff, source_task: Task) -> TaskResult:
        """Transfer: fully hand over to another agent."""
        source_task.description = handoff.message or source_task.description
        source_task.status = TaskStatus.SUBMITTED  # Reset status for new agent
        result = await self.run_task(source_task, agent_id=handoff.target_agent_id)
        handoff.complete(result.content)
        return result

    async def _broadcast(self, handoff: Handoff, source_task: Task) -> TaskResult:
        """Broadcast: fan out to multiple agents, collect results."""
        coros = []
        for agent_id in handoff.target_agent_ids:
            sub_task = Task(
                name=f"broadcast_{source_task.name}_{agent_id}",
                description=handoff.message,
                input_data=handoff.context_data,
            )
            coros.append(self.run_task(sub_task, agent_id=agent_id))

        results = await asyncio.gather(*coros, return_exceptions=True)

        # Collect successful results
        successful = {}
        for agent_id, result in zip(handoff.target_agent_ids, results):
            if isinstance(result, TaskResult):
                successful[agent_id] = result.content
                handoff.results[agent_id] = result.content
            elif isinstance(result, Exception):
                handoff.results[agent_id] = f"ERROR: {result}"

        # Pick the best result (simple: longest non-error response)
        best_content = max(successful.values(), key=len) if successful else "All agents failed"
        handoff.complete(best_content)

        return TaskResult(
            content=best_content,
            data={"broadcast_results": handoff.results},
            metadata={"broadcast_count": len(handoff.target_agent_ids)},
        )

    # ---- Resume from checkpoint ----

    async def resume_task(
        self,
        task_id: str,
        agent_id: str | None = None,
    ) -> TaskResult:
        """
        Resume a task from its latest checkpoint.

        Loads the checkpoint, restores state, and continues execution.
        """
        checkpoint = await self.checkpoint_store.get_latest(task_id)
        if not checkpoint:
            raise ValueError(f"No checkpoint found for task {task_id}")

        # Reconstruct task from checkpoint
        task = Task(
            id=task_id,
            name=checkpoint.metadata.get("task_name", ""),
            description=checkpoint.metadata.get("description", ""),
            status=TaskStatus.SUBMITTED,  # Reset for re-execution
            input_data=checkpoint.state_data,
        )

        # Restore conversation
        for msg_data in checkpoint.conversation_messages:
            msg = Message.model_validate(msg_data)
            task.conversation.append(msg)

        resolved_agent_id = agent_id or checkpoint.agent_id
        logger.info(
            f"Resuming task {task_id} from checkpoint {checkpoint.id} "
            f"(iteration={checkpoint.iteration})"
        )

        return await self.run_task(task, agent_id=resolved_agent_id)

    # ---- Internal helpers ----

    def _resolve_agent(self, task: Task, agent_id: str | None) -> Agent | None:
        """Find the agent to execute a task."""
        if agent_id:
            return self.discovery.get_agent(agent_id)

        # Auto-select based on capabilities
        if task.required_capabilities:
            best_id = self.discovery.find_best_agent(task.required_capabilities)
            if best_id:
                return self.discovery.get_agent(best_id)

        # Fallback: pick any available agent
        cards = self.discovery.list_agents()
        for card in cards:
            agent = self.discovery.get_agent(card.agent_id)
            if agent:
                return agent
        return None

    def _make_fallback_executor(self, task: Task):
        """Create a fallback executor factory for the policy chain."""
        def factory(agent_id: str):
            async def executor(t: Task) -> TaskResult:
                agent = self.discovery.get_agent(agent_id)
                if not agent:
                    raise RuntimeError(f"Fallback agent {agent_id} not found")
                return await agent.run(t)
            return executor
        return factory

    async def _save_checkpoint(self, task: Task, agent_id: str) -> str:
        """Save a checkpoint for the current task state."""
        checkpoint = Checkpoint(
            task_id=task.id,
            agent_id=agent_id,
            state_data=task.input_data,
            conversation_messages=[
                m.model_dump() for m in task.conversation.messages
            ],
            token_usage=task.token_usage.model_dump(),
            status=task.status.value,
            metadata={
                "task_name": task.name,
                "description": task.description,
            },
        )
        return await self.checkpoint_store.save(checkpoint)

    def _record_quality(
        self,
        task: Task,
        agent_id: str,
        result: TaskResult,
    ) -> None:
        """Record task quality for the discovery service's feedback loop."""
        quality_score = result.quality_score if result.quality_score is not None else 0.5
        record = QualityRecord(
            agent_id=agent_id,
            task_id=task.id,
            task_type=task.task_type,
            score=quality_score,
            tokens_used=task.token_usage.total_tokens,
            cost_usd=task.token_usage.cost_usd,
            latency_seconds=task.elapsed_seconds or 0.0,
        )
        self.discovery.record_quality(record)
