"""Team — multi-agent orchestration for Logos ADK."""
from __future__ import annotations

import logging
from typing import Any, AsyncIterator
from uuid import uuid4

from logos_adk.a2a import A2ARequest, A2AResponse
from logos_adk.errors import MaxRoundsExceeded
from logos_adk.handler import collect_handlers
from logos_adk.mailbox import Mailbox
from logos_adk.runnable import Runnable
from logos_adk.team_config import TeamConfig
from logos_adk.tools import Tool
from logos_adk.topology import resolve_topology
from logos_adk.types import Response, StreamChunk, Usage
from logos_adk.ui_events import make_ui_event

from logos_adk.team._budgeting import BudgetMixin
from logos_adk.team._checkpoint import CheckpointMixin
from logos_adk.team._crew import CrewMixin
from logos_adk.team._event_bus import EventBusMixin
from logos_adk.team._messaging import MessagingMixin
from logos_adk.team._types import (
    AgentLike,
    CampaignState as CampaignState,
    _PausedCrewRun as _PausedCrewRun,
    _TeamEventSubscription as _TeamEventSubscription,
)


class Team(
    CrewMixin,
    CheckpointMixin,
    EventBusMixin,
    BudgetMixin,
    MessagingMixin,
    Runnable,
):
    """A group of agents that collaborate to complete tasks."""

    def __init__(
        self,
        name: str,
        *,
        agents: list[AgentLike] | None = None,
        orchestrator: str | None = None,
        topology: str | dict[str, list[str]] | None = None,
        max_rounds: int = 10,
        supervisor: Any | None = None,
        process: str | None = None,
        manager_role: str | None = None,
        checkpoint_store: Any | None = None,
        event_bus: Any | None = None,
        topic_policy: dict[str, list[str]] | None = None,
    ) -> None:
        agents = agents or []
        self._config = TeamConfig(
            name=name,
            agents=agents,
            orchestrator=orchestrator,
            topology=topology,
            max_rounds=max_rounds,
            supervisor=supervisor,
            process=process,
            manager_role=manager_role,
            topic_policy=topic_policy,
        )
        self._agents: dict[str, AgentLike] = {}
        self._mailboxes: dict[str, Mailbox] = {}
        self._pending_a2a_responses: dict[str, A2AResponse] = {}
        self._paused_crew_runs: dict[str, _PausedCrewRun] = {}
        self._checkpoint_store = checkpoint_store
        self._event_bus = event_bus
        self._topic_subscriptions: list[_TeamEventSubscription] = []
        self._team_budget_usd: float | None = None
        self._team_budget_tokens: int | None = None
        self._task_budget_usd: float | None = None
        self._task_budget_tokens: int | None = None
        self._budget_threshold_ratio: float = 0.8
        self._fail_on_budget_exhausted: bool = False

        # Register agents and detect duplicates
        for agent in agents:
            if agent.name in self._agents:
                raise ValueError(
                    f"duplicate agent name '{agent.name}' in team '{name}'"
                )
            self._agents[agent.name] = agent
            self._mailboxes[agent.name] = Mailbox()
            # Set backref for Agent instances; nested Teams don't have _team
            if hasattr(agent, '_team'):
                agent._team = self

        # Resolve orchestrator
        if orchestrator is not None and orchestrator not in self._agents:
            raise ValueError(
                f"Orchestrator '{orchestrator}' not in agents: "
                f"{list(self._agents.keys())}"
            )
        self._orchestrator_name: str | None = orchestrator or (
            agents[0].name if agents else None
        )

        # Collect @handler methods per agent
        self._handlers: dict[str, dict[type, Any]] = {}
        for agent in agents:
            self._handlers[agent.name] = collect_handlers(agent)

        # Resolve topology
        if agents:
            self._topology = resolve_topology(
                topology, list(self._agents.keys())
            )
        else:
            self._topology: dict[str, set[str]] = {}
        self._topic_policy = self._resolve_topic_policy(topic_policy)
        self._restore_persistent_tools()
        self._register_builtin_scheduled_task_subscriptions()
        for agent in agents:
            self._register_agent_topic_subscriptions(agent)

    @property
    def name(self) -> str:
        return self._config.name

    def _restore_persistent_tools(self) -> None:
        try:
            from logos_adk.tool_synthesis import GeneratedToolStore

            GeneratedToolStore().attach_tools_for_team(self)
        except (ImportError, AttributeError):
            # tool_synthesis is an optional module; skip if unavailable
            return
        except Exception:
            logging.getLogger(__name__).debug(
                "Failed to restore persistent tools for team '%s'",
                self._config.name,
                exc_info=True,
            )
            return

    def artifact_sandbox(self, *, root: str | None = None) -> Team:
        """Attach the shared artifact sandbox toolkit to every team member."""
        for agent in self._agents.values():
            if hasattr(agent, "artifact_sandbox"):
                agent.artifact_sandbox(root=root)
        return self

    def temporal_autonomy(self) -> Team:
        """Attach temporal scheduling tools to every team member."""
        for agent in self._agents.values():
            if hasattr(agent, "temporal_autonomy"):
                agent.temporal_autonomy()
        return self

    @classmethod
    def from_definition(cls, config: TeamConfig) -> Team:
        """Build a Team from a typed TeamConfig."""
        from logos_adk.agent import Agent
        from logos_adk.config_loader import build_agent_from_definition
        from logos_adk.errors import ConfigError

        agents = [build_agent_from_definition(agent_config, default_cls=Agent) for agent_config in config.agents]
        event_bus = cls._build_event_bus_from_config(config.event_bus)
        team = cls(
            config.name,
            agents=agents,
            orchestrator=config.orchestrator,
            topology=config.topology,
            max_rounds=config.max_rounds,
            supervisor=config.supervisor,
            process=config.process,
            manager_role=config.manager_role,
            event_bus=event_bus,
            topic_policy=config.topic_policy,
        )
        for subscription in config.subscriptions:
            try:
                team.subscribe(
                    subscription.agent_name,
                    subscription.topic,
                    handler=subscription.handler,
                    publish_to=subscription.publish_to,
                    name=subscription.name,
                )
            except Exception as exc:
                raise ConfigError(
                    f"Failed to register subscription '{subscription.agent_name}:{subscription.topic}': {exc}"
                ) from exc
        return team

    @classmethod
    def from_config(cls, path: str) -> Team:
        """Load a Team from a YAML config file."""
        from logos_adk.team_config_parser import load_team_config

        team = cls.from_definition(load_team_config(path))
        for agent in team._agents.values():
            if hasattr(agent, "_config_path"):
                agent._config_path = path
        return team

    def get_agent(self, name: str) -> AgentLike | None:
        """Get a registered agent by name."""
        return self._agents.get(name)

    def get_agents_by_role(self, role: str) -> list[AgentLike]:
        """Return all agents assigned to a Crew-style role."""
        return [
            agent
            for agent in self._agents.values()
            if getattr(getattr(agent, "_config", None), "role", None) == role
        ]

    def get_agent_by_role(self, role: str) -> AgentLike | None:
        """Return the first agent assigned to a Crew-style role."""
        matches = self.get_agents_by_role(role)
        return matches[0] if matches else None

    def roster(self) -> list[dict[str, str]]:
        """Return team roster metadata for role-aware orchestration."""
        roster: list[dict[str, str]] = []
        for agent in self._agents.values():
            profile = getattr(agent, "crew_profile", {})
            if not profile:
                continue
            roster.append({"name": agent.name, **profile})
        return roster

    # -- execution -----------------------------------------------------------

    async def kickoff(
        self,
        prompt: str,
        *,
        role: str | None = None,
        session_id: str | None = None,
        **kwargs: Any,
    ) -> Response:
        """Crew-style kickoff that can route directly by role."""
        effective_kwargs = self._apply_budget_kwargs(kwargs)
        if role is not None:
            target = self.get_agent_by_role(role)
            if target is None:
                raise ValueError(f"Role '{role}' is not assigned to any agent in team '{self.name}'")
            return await target.run(prompt, session_id=session_id, **effective_kwargs)
        if self.roster():
            prompt = self._role_aware_prompt(prompt)
        return await self.run(prompt, session_id=session_id, **effective_kwargs)

    async def delegate(
        self,
        task: str,
        *,
        role: str | None = None,
        to_agent: str | None = None,
        session_id: str | None = None,
        **kwargs: Any,
    ) -> Response:
        """Delegate a task directly to an agent or to the first agent with a given role."""
        effective_kwargs = self._apply_budget_kwargs(kwargs)
        if to_agent is None:
            if role is None:
                raise ValueError("delegate() requires either role or to_agent")
            agent = self.get_agent_by_role(role)
            if agent is None:
                raise ValueError(f"Role '{role}' is not assigned to any agent in team '{self.name}'")
        else:
            agent = self.get_agent(to_agent)
            if agent is None:
                raise ValueError(f"Agent '{to_agent}' is not registered in team '{self.name}'")
        return await agent.run(task, session_id=session_id, **effective_kwargs)

    async def run(
        self, prompt: str, *, session_id: str | None = None, **kwargs: Any
    ) -> Response:
        """Run the team: send prompt to orchestrator and return its response."""
        if not self._agents:
            raise ValueError(f"Team '{self.name}' has no agents")

        orchestrator = self._agents[self._orchestrator_name]
        response = await orchestrator.run(
            prompt,
            session_id=session_id,
            **self._apply_budget_kwargs(kwargs),
        )
        return response

    def _has_pending_messages(self) -> bool:
        """Check if any agent mailbox has pending messages."""
        return any(not mb.empty() for mb in self._mailboxes.values())

    async def run_rounds(self) -> Response:
        """Process mailbox messages in rounds until all mailboxes are empty.

        Raises:
            MaxRoundsExceeded: If max_rounds is exceeded and messages remain.
        """
        last_response: Response | None = None

        for _ in range(self._config.max_rounds):
            any_activity = False

            for agent_name in list(self._agents.keys()):
                mailbox = self._mailboxes[agent_name]
                messages = mailbox.drain()
                if not messages:
                    continue
                any_activity = True
                agent = self._agents[agent_name]
                handlers = self._handlers.get(agent_name, {})
                for msg in messages:
                    if isinstance(msg, A2AResponse):
                        self._pending_a2a_responses[msg.request_id] = msg
                        last_response = Response(
                            content=msg.content,
                            tool_calls=[],
                            usage=Usage(0, 0),
                            raw={"a2a_response": self._serialize_a2a_response(msg)},
                        )
                        continue
                    if isinstance(msg, A2ARequest):
                        handler = handlers.get(A2ARequest)
                        if handler is not None:
                            result = handler(msg)
                            if hasattr(result, '__await__'):
                                await result
                        else:
                            prompt = self._build_a2a_prompt(msg)
                            response = await agent.run(prompt)
                            if msg.expects_response:
                                await self.send(
                                    agent_name,
                                    msg.from_agent,
                                    A2AResponse(
                                        from_agent=agent_name,
                                        to_agent=msg.from_agent,
                                        content=response.content,
                                        conversation_id=msg.conversation_id,
                                        metadata={"task_id": msg.task_id},
                                        request_id=msg.task_id,
                                        status="completed",
                                    ),
                                )
                            last_response = response
                        continue
                    # Dispatch to @handler if message type matches
                    handler = handlers.get(type(msg))
                    if handler is not None:
                        result = handler(msg)
                        if hasattr(result, '__await__'):
                            await result
                    else:
                        content = msg if isinstance(msg, str) else str(msg)
                        last_response = await agent.run(content)

            event_response, processed_events = await self._run_event_round()
            if processed_events > 0:
                any_activity = True
                if event_response is not None:
                    last_response = event_response

            if not any_activity:
                break

        if self._has_pending_messages():
            raise MaxRoundsExceeded(
                f"Team '{self.name}' exceeded max_rounds={self._config.max_rounds}"
            )

        if last_response is None:
            last_response = Response(
                content="", tool_calls=[], usage=Usage(0, 0), raw={}
            )
        return last_response

    # -- streaming -----------------------------------------------------------

    async def stream(
        self, prompt: str, *, session_id: str | None = None, **kwargs: Any
    ) -> AsyncIterator[StreamChunk]:
        """Stream the team execution."""
        if not self._agents:
            raise ValueError(f"Team '{self.name}' has no agents")

        orchestrator = self._agents[self._orchestrator_name]
        async for chunk in orchestrator.stream(prompt, session_id=session_id, **kwargs):
            yield chunk

    async def stream_events(
        self,
        prompt: str,
        *,
        session_id: str | None = None,
        correlation_id: str | None = None,
        **kwargs: Any,
    ) -> AsyncIterator[dict[str, Any]]:
        run_id = uuid4().hex
        sequence = 0
        trace_id = kwargs.get("_trace_id")
        effective_correlation_id = correlation_id or session_id or run_id
        yield make_ui_event(
            event="graph_started",
            event_type="team",
            graph_name=self.name,
            graph_run_id=run_id,
            sequence=sequence,
            session_id=session_id,
            correlation_id=effective_correlation_id,
            trace_id=trace_id,
            node_name=self.name,
            status="running",
        )
        sequence += 1
        orchestrator = self._agents[self._orchestrator_name]
        final_event: dict[str, Any] | None = None
        async for event in orchestrator.stream_events(
            prompt,
            session_id=session_id,
            correlation_id=effective_correlation_id,
            **kwargs,
        ):
            payload = dict(event)
            payload.setdefault("metadata", {})
            payload["metadata"]["team_name"] = self.name
            if payload.get("event") == "graph_started":
                continue
            if payload.get("event") == "graph_completed":
                final_event = payload
                continue
            yield payload
        yield make_ui_event(
            event="graph_completed",
            event_type="team",
            graph_name=self.name,
            graph_run_id=run_id,
            sequence=sequence,
            session_id=session_id,
            correlation_id=effective_correlation_id,
            trace_id=trace_id,
            node_name=self.name,
            status=(final_event or {}).get("status") or "completed",
            output=(final_event or {}).get("output"),
            metadata=(final_event or {}).get("metadata"),
        )

    # -- tool conversion -----------------------------------------------------

    def as_tool(
        self, *, name: str | None = None, description: str | None = None
    ) -> Tool:
        """Convert this team into a Tool that can be used by other agents."""
        tool_name = name or self.name
        tool_desc = description or f"Delegate to the '{self.name}' team"
        team = self

        async def _team_tool(prompt: str) -> str:
            response = await team.run(prompt)
            return response.content

        return Tool(fn=_team_tool, name=tool_name, description=tool_desc)

    def _role_aware_prompt(self, prompt: str) -> str:
        roster_lines = []
        for entry in self.roster():
            line = f"- {entry['name']}: role={entry.get('role', '')}"
            if entry.get("goal"):
                line += f"; goal={entry['goal']}"
            if entry.get("backstory"):
                line += f"; backstory={entry['backstory']}"
            roster_lines.append(line)
        return (
            "Team roster:\n"
            + "\n".join(roster_lines)
            + "\n\nTask:\n"
            + prompt
        )
