from __future__ import annotations

from pydantic import BaseModel, Field

from nexus.resilience.errors import OrchestratorError


class AgentProfile(BaseModel):
    """Agent capability and performance profile."""

    agent_id: str
    role: str
    capabilities: list[str] = Field(default_factory=list)
    avg_quality_score: float = 0.5
    avg_token_efficiency: float = 1.0
    task_count: int = 0
    specialty_tags: list[str] = Field(default_factory=list)


class RoleAssignor:
    """Dynamic role assignment based on capability and outcomes."""

    def __init__(self) -> None:
        self._profiles: dict[str, AgentProfile] = {}

    def register_agent(self, profile: AgentProfile) -> None:
        """Register or update an agent profile."""
        self._profiles[profile.agent_id] = profile

    def update_performance(self, agent_id: str, quality_score: float, token_efficiency: float) -> None:
        """Update running performance averages."""
        profile = self._profiles.get(agent_id)
        if profile is None:
            raise OrchestratorError(f"Unknown agent profile: {agent_id}")
        count = profile.task_count
        profile.avg_quality_score = ((profile.avg_quality_score * count) + quality_score) / (count + 1)
        profile.avg_token_efficiency = ((profile.avg_token_efficiency * count) + token_efficiency) / (count + 1)
        profile.task_count += 1

    def assign(
        self,
        task_tags: list[str],
        required_capabilities: list[str] | None = None,
        token_budget: int | None = None,
    ) -> str:
        """Return best matching agent id under weighted scoring."""
        if not self._profiles:
            raise OrchestratorError("No agents registered")
        req_caps = required_capabilities or []
        best_id = ""
        best_score = -1.0
        best_quality = -1.0

        for profile in self._profiles.values():
            if req_caps:
                matches = len(set(req_caps).intersection(profile.capabilities))
                capability_score = matches / len(req_caps)
            else:
                capability_score = 1.0
            specialty_score = len(set(task_tags).intersection(profile.specialty_tags)) / max(len(task_tags), 1)
            performance_weight = (profile.avg_quality_score * 0.6) + (profile.avg_token_efficiency * 0.4)
            total_score = (0.4 * capability_score) + (0.3 * specialty_score) + (0.3 * performance_weight)

            if token_budget is not None and token_budget < 1000:
                total_score *= 0.95

            if (total_score > best_score) or (total_score == best_score and profile.avg_quality_score > best_quality):
                best_score = total_score
                best_quality = profile.avg_quality_score
                best_id = profile.agent_id

        if not best_id:
            raise OrchestratorError("No suitable agent assignment found")
        return best_id
