"""
Agent Discovery Service — find and match agents to tasks.

Design notes:
  - Local registry: agents register on startup, look up by capability match
  - Remote discovery: fetch AgentCards from well-known URIs (A2A compatible)
  - Semantic matching: compute cosine similarity between task requirements
    and agent capabilities to find the best match
  - ELO-like rating: agents build a track record that influences future matching
"""

from __future__ import annotations

import logging
from datetime import datetime

from pydantic import BaseModel, Field

from nexus.core.agent import Agent
from nexus.types import AgentCapability, AgentCard, QualityRecord

logger = logging.getLogger("nexus.protocol.discovery")


class AgentScore(BaseModel):
    """Tracks an agent's cumulative performance score (ELO-inspired)."""

    agent_id: str
    rating: float = 1000.0  # Initial ELO-like rating
    tasks_completed: int = 0
    tasks_failed: int = 0
    avg_quality: float = 0.0
    avg_latency: float = 0.0
    total_cost: float = 0.0
    last_active: datetime = Field(default_factory=datetime.utcnow)

    def update(self, record: QualityRecord) -> None:
        """Update the agent's rating based on a new quality record."""
        self.tasks_completed += 1
        # Running average for quality
        n = self.tasks_completed
        self.avg_quality = ((n - 1) * self.avg_quality + record.score) / n
        self.avg_latency = ((n - 1) * self.avg_latency + record.latency_seconds) / n
        self.total_cost += record.cost_usd
        self.last_active = record.timestamp

        # ELO-like update: higher quality → higher rating
        k_factor = 32.0
        expected = 1.0 / (1.0 + 10 ** ((1000 - self.rating) / 400))
        self.rating += k_factor * (record.score - expected)

    def record_failure(self) -> None:
        """Record a task failure (decreases rating)."""
        self.tasks_failed += 1
        self.rating = max(0, self.rating - 16)  # Penalty


class DiscoveryService:
    """
    Agent registry and discovery service.

    Manages the pool of available agents and matches them to tasks
    based on capability similarity and historical performance.

    Usage:
        discovery = DiscoveryService()

        # Register local agents
        discovery.register(agent_a)
        discovery.register(agent_b)

        # Find the best agent for a task
        candidates = discovery.find_agents(
            required=AgentCapability(coding=0.8, reasoning=0.7),
            top_k=3,
        )

        # Record performance feedback
        discovery.record_quality(QualityRecord(
            agent_id=agent_a.id,
            task_id="task-123",
            score=0.92,
        ))
    """

    def __init__(self) -> None:
        self._agents: dict[str, Agent] = {}
        self._cards: dict[str, AgentCard] = {}
        self._scores: dict[str, AgentScore] = {}
        self._quality_history: list[QualityRecord] = []

    # ---- Registration ----

    def register(self, agent: Agent) -> AgentCard:
        """Register a local agent and return its AgentCard.

        Idempotent: calling multiple times with the same agent is safe
        and will only log once.
        """
        is_new = agent.id not in self._agents

        self._agents[agent.id] = agent
        card = agent.to_agent_card()
        self._cards[agent.id] = card

        if agent.id not in self._scores:
            self._scores[agent.id] = AgentScore(agent_id=agent.id)

        if is_new:
            logger.info(f"Registered agent '{agent.name}' (id={agent.id})")
        return card

    def register_remote(self, card: AgentCard) -> None:
        """Register a remote agent via its AgentCard (discovered via A2A)."""
        self._cards[card.agent_id] = card
        if card.agent_id not in self._scores:
            self._scores[card.agent_id] = AgentScore(agent_id=card.agent_id)
        logger.info(f"Registered remote agent '{card.name}' (id={card.agent_id})")

    def unregister(self, agent_id: str) -> None:
        """Remove an agent from the registry."""
        self._agents.pop(agent_id, None)
        self._cards.pop(agent_id, None)
        logger.info(f"Unregistered agent {agent_id}")

    # ---- Discovery ----

    def get_agent(self, agent_id: str) -> Agent | None:
        """Look up a local agent by ID."""
        return self._agents.get(agent_id)

    def get_card(self, agent_id: str) -> AgentCard | None:
        """Look up an agent's card by ID."""
        return self._cards.get(agent_id)

    def list_agents(self) -> list[AgentCard]:
        """List all registered agent cards."""
        return list(self._cards.values())

    def find_agents(
        self,
        required: AgentCapability,
        top_k: int = 5,
        min_rating: float = 0.0,
        exclude: set[str] | None = None,
    ) -> list[tuple[str, float]]:
        """
        Find the best agents for a given capability requirement.

        Returns a list of (agent_id, match_score) tuples, sorted by score descending.
        The match_score combines capability similarity and historical performance.

        Args:
            required: The capability vector describing what the task needs.
            top_k: Maximum number of candidates to return.
            min_rating: Minimum ELO rating to be considered.
            exclude: Agent IDs to exclude from results.
        """
        exclude = exclude or set()
        candidates: list[tuple[str, float]] = []

        for agent_id, card in self._cards.items():
            if agent_id in exclude:
                continue

            score_record = self._scores.get(agent_id)
            if score_record and score_record.rating < min_rating:
                continue

            # Capability similarity (0.0 - 1.0)
            cap_score = card.capabilities.cosine_similarity(required)

            # Historical performance weight (0.0 - 1.0, normalized)
            perf_score = 0.5  # Default for agents with no history
            if score_record and score_record.tasks_completed > 0:
                perf_score = score_record.rating / 2000.0  # Normalize to ~0-1 range

            # Combined score: 60% capability match, 40% track record
            combined = 0.6 * cap_score + 0.4 * perf_score
            candidates.append((agent_id, combined))

        # Sort by score descending
        candidates.sort(key=lambda x: x[1], reverse=True)
        return candidates[:top_k]

    def find_best_agent(
        self,
        required: AgentCapability,
        exclude: set[str] | None = None,
    ) -> str | None:
        """Find the single best agent for a capability requirement."""
        results = self.find_agents(required, top_k=1, exclude=exclude)
        return results[0][0] if results else None

    # ---- Feedback loop ----

    def record_quality(self, record: QualityRecord) -> None:
        """Record a quality signal to update agent ratings."""
        self._quality_history.append(record)
        score = self._scores.get(record.agent_id)
        if score:
            score.update(record)
            logger.debug(
                f"Agent {record.agent_id} quality recorded: "
                f"score={record.score:.2f}, new_rating={score.rating:.0f}"
            )

    def record_failure(self, agent_id: str, task_id: str) -> None:
        """Record a task failure for an agent."""
        score = self._scores.get(agent_id)
        if score:
            score.record_failure()
            logger.debug(
                f"Agent {agent_id} failure recorded, "
                f"new_rating={score.rating:.0f}"
            )

    def get_score(self, agent_id: str) -> AgentScore | None:
        """Get an agent's performance score."""
        return self._scores.get(agent_id)

    def get_leaderboard(self) -> list[AgentScore]:
        """Get all agents sorted by rating (descending)."""
        return sorted(
            self._scores.values(),
            key=lambda s: s.rating,
            reverse=True,
        )

    # ---- Remote discovery (A2A) ----

    async def discover_remote(self, base_url: str) -> AgentCard | None:
        """
        Discover a remote agent by fetching its AgentCard.

        Uses A2A's well-known URI convention.
        """
        from nexus.protocol.a2a import A2AClient

        client = A2AClient()
        try:
            card = await client.discover(base_url)
            if card:
                self.register_remote(card)
            return card
        finally:
            await client.close()

    @property
    def agent_count(self) -> int:
        return len(self._cards)

    @property
    def local_agent_count(self) -> int:
        return len(self._agents)
