"""
AdaptiveOrchestrator — self-organizing topology mode (the core innovation).

Design notes:
  - Given a set of tasks and a pool of agents, automatically finds the
    optimal assignment based on: capability matching + historical performance + constraints
  - Supports re-planning: if a task fails or quality is low, reassign to
    a different agent or decompose into sub-tasks
  - Uses a Contract Net-inspired protocol: tasks are "announced" and agents
    "bid" based on their capabilities and current load
  - Tracks performance to continuously improve future assignments
"""

from __future__ import annotations

import asyncio
import logging
from dataclasses import dataclass, field

from nexus.core.agent import Agent
from nexus.core.task import Task, TaskGraph, TaskResult, TaskStatus
from nexus.protocol.discovery import DiscoveryService
from nexus.types import AgentCapability, QualityRecord

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


@dataclass
class Assignment:
    """Records an agent-task assignment with metadata."""
    task_id: str
    agent_id: str
    match_score: float = 0.0
    reason: str = ""


@dataclass
class PlanResult:
    """The output of the adaptive planner — a set of assignments."""
    assignments: list[Assignment] = field(default_factory=list)
    unassigned_tasks: list[str] = field(default_factory=list)

    def get_agent_for_task(self, task_id: str) -> str | None:
        for a in self.assignments:
            if a.task_id == task_id:
                return a.agent_id
        return None


class AdaptiveOrchestrator:
    """
    Self-organizing topology — agents are dynamically matched to tasks
    based on capabilities, constraints, and historical performance.

    This is the ATMAS core: the system discovers the optimal topology
    at runtime rather than having it pre-defined.

    Usage:
        orchestrator = AdaptiveOrchestrator(discovery=discovery)

        # Add agents to the pool
        orchestrator.add_agent(researcher)
        orchestrator.add_agent(coder)
        orchestrator.add_agent(writer)

        # Plan and execute
        plan = orchestrator.plan(task_graph)
        results = await orchestrator.execute(task_graph, plan)

        # Or let it auto-plan and execute
        results = await orchestrator.run(task_graph)
    """

    def __init__(
        self,
        discovery: DiscoveryService | None = None,
        quality_threshold: float = 0.6,
        max_replanning_rounds: int = 3,
    ) -> None:
        self.discovery = discovery or DiscoveryService()
        self.quality_threshold = quality_threshold
        self.max_replanning_rounds = max_replanning_rounds
        self._agents: dict[str, Agent] = {}

    def add_agent(self, agent: Agent) -> None:
        """Add an agent to the available pool."""
        self._agents[agent.id] = agent
        self.discovery.register(agent)

    # ---- Planning ----

    def plan(self, graph: TaskGraph) -> PlanResult:
        """
        Generate an optimal assignment plan for a task graph.

        For each task, find the best agent based on:
        - Capability vector similarity (60% weight)
        - Historical performance rating (40% weight)
        - Constraint satisfaction (hard filter)

        This is a greedy algorithm. For more optimal solutions, consider
        using Hungarian algorithm or constraint satisfaction solvers.
        """
        result = PlanResult()
        assigned_agents: dict[str, int] = {}  # agent_id → task count (for load balancing)

        # Sort tasks by priority (tasks with more dependencies first)
        sorted_tasks = sorted(
            graph.tasks.values(),
            key=lambda t: len(t.depends_on),
            reverse=True,
        )

        for task in sorted_tasks:
            agent_id = self._select_best_agent(task, assigned_agents)

            if agent_id:
                result.assignments.append(
                    Assignment(
                        task_id=task.id,
                        agent_id=agent_id,
                        match_score=self._compute_match_score(task, agent_id),
                        reason="Best match by capability+performance",
                    )
                )
                assigned_agents[agent_id] = assigned_agents.get(agent_id, 0) + 1
            else:
                result.unassigned_tasks.append(task.id)
                logger.warning(f"No suitable agent for task '{task.name}'")

        logger.info(
            f"Adaptive plan: {len(result.assignments)} assigned, "
            f"{len(result.unassigned_tasks)} unassigned"
        )
        return result

    def _select_best_agent(
        self,
        task: Task,
        current_load: dict[str, int],
    ) -> str | None:
        """Select the best agent for a task considering capabilities, performance, and load."""
        if not self._agents:
            logger.warning("No agents registered in adaptive orchestrator")
            return None

        required = task.required_capabilities or AgentCapability()
        candidates = self.discovery.find_agents(required, top_k=10)

        if candidates:
            # Re-rank considering current load (prefer less-loaded agents)
            scored = []
            for agent_id, base_score in candidates:
                load = current_load.get(agent_id, 0)
                load_penalty = load * 0.1  # Each assigned task reduces score by 0.1
                final_score = base_score - load_penalty

                # Check constraint satisfaction
                if not self._check_constraints(task, agent_id):
                    continue

                scored.append((agent_id, final_score))

            if scored:
                scored.sort(key=lambda x: x[1], reverse=True)
                winner_id = scored[0][0]
                winner_agent = self._agents.get(winner_id)
                logger.debug(
                    f"Selected agent '{winner_agent.name if winner_agent else winner_id}' "
                    f"(score={scored[0][1]:.3f}) for task '{task.name}'"
                )
                return winner_id

        # Fallback: any available agent (better than returning None)
        fallback_id = next(iter(self._agents.keys()))
        fallback_agent = self._agents[fallback_id]
        logger.info(
            f"No scored match, falling back to agent '{fallback_agent.name}' "
            f"for task '{task.name}'"
        )
        return fallback_id

    def _compute_match_score(self, task: Task, agent_id: str) -> float:
        """Compute the match score between a task and an agent."""
        card = self.discovery.get_card(agent_id)
        if not card:
            return 0.0
        required = task.required_capabilities or AgentCapability()
        return card.capabilities.cosine_similarity(required)

    def _check_constraints(self, task: Task, agent_id: str) -> bool:
        """Check if an agent can satisfy a task's constraints."""
        # For now, all agents pass constraint checks.
        # In production, this would check:
        # - Model cost vs task budget
        # - Expected latency vs task deadline
        # - Required tools vs agent tools
        return True

    # ---- Execution ----

    async def execute(
        self,
        graph: TaskGraph,
        plan: PlanResult,
    ) -> dict[str, TaskResult]:
        """Execute a task graph following the plan."""
        results: dict[str, TaskResult] = {}

        while not graph.is_complete:
            ready_tasks = graph.get_ready_tasks()
            if not ready_tasks:
                remaining = [
                    t for t in graph.tasks.values() if not t.status.is_terminal
                ]
                if remaining:
                    raise RuntimeError("Task graph deadlock")
                break

            # Execute ready tasks in parallel
            coros = []
            for task in ready_tasks:
                agent_id = plan.get_agent_for_task(task.id)
                coros.append(self._execute_task(task, agent_id, results))

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

        return results

    async def _execute_task(
        self,
        task: Task,
        agent_id: str | None,
        results: dict[str, TaskResult],
    ) -> None:
        """Execute a single task with the assigned agent."""
        agent = self._agents.get(agent_id) if agent_id else None
        if not agent:
            # Fallback: try to find any agent
            if self._agents:
                agent = next(iter(self._agents.values()))
            else:
                task.fail("No agent available")
                return

        # 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

        try:
            result = await agent.run(task)
            results[task.id] = result

            # Record quality for future planning
            quality = result.quality_score or 0.5
            self.discovery.record_quality(
                QualityRecord(
                    agent_id=agent.id,
                    task_id=task.id,
                    task_type=task.task_type,
                    score=quality,
                    tokens_used=task.token_usage.total_tokens,
                    latency_seconds=task.elapsed_seconds or 0.0,
                )
            )
        except Exception as exc:
            logger.error(f"Task '{task.name}' failed with agent '{agent.name}': {exc}")
            task.fail(str(exc))
            self.discovery.record_failure(agent.id, task.id)

    # ---- Run with replanning ----

    async def run(self, graph: TaskGraph) -> dict[str, TaskResult]:
        """
        Plan, execute, and re-plan if needed.

        If tasks fail or produce low-quality results, the orchestrator
        re-plans with different agent assignments and retries.
        """
        for round_num in range(1, self.max_replanning_rounds + 1):
            logger.info(f"Adaptive execution round {round_num}")

            plan = self.plan(graph)
            results = await self.execute(graph, plan)

            # Check if re-planning is needed
            needs_replan = False
            for task_id, task in graph.tasks.items():
                if task.status == TaskStatus.FAILED and task.can_retry():
                    task.status = TaskStatus.SUBMITTED  # Reset for retry
                    task.retry_count += 1
                    needs_replan = True
                    logger.info(
                        f"Task '{task.name}' failed, scheduling retry "
                        f"({task.retry_count}/{task.max_retries})"
                    )

                # Check quality threshold
                result = results.get(task_id)
                if result and result.quality_score is not None:
                    if result.quality_score < self.quality_threshold and task.can_retry():
                        task.status = TaskStatus.SUBMITTED
                        task.retry_count += 1
                        needs_replan = True
                        logger.info(
                            f"Task '{task.name}' quality {result.quality_score:.2f} "
                            f"< threshold {self.quality_threshold}, scheduling retry"
                        )

            if not needs_replan:
                break

        return results
