"""Application service for coordinated team runs."""

from __future__ import annotations

from collections.abc import Callable
from typing import TYPE_CHECKING

from beaver.coordinator import ExecutionGraph, ExecutionNode, LocalAgentRunner, TeamGraphScheduler, TeamRunResult
from beaver.engine import AgentLoop
from beaver.engine.providers import ProviderBundle

if TYPE_CHECKING:
    from beaver.engine.context import SkillContext


class TeamService:
    """Internal service for Beaver-native multi-agent execution."""

    def __init__(self, loop: AgentLoop) -> None:
        self.loop = loop
        self.runner = LocalAgentRunner(loop)
        self.scheduler = TeamGraphScheduler(self.runner)

    async def run_team(
        self,
        graph: ExecutionGraph,
        *,
        parent_task_id: str | None,
        parent_session_id: str,
        parent_run_id: str | None = None,
        provider_bundle: ProviderBundle | None = None,
        provider_bundle_factory: Callable[[ExecutionNode], ProviderBundle | None] | None = None,
        inherited_pinned_skills: list[str] | None = None,
        inherited_pinned_skill_contexts: list["SkillContext"] | None = None,
        allow_candidate_generation: bool = False,
    ) -> TeamRunResult:
        """Run a team graph inside the parent task context."""

        self._validate_parent_task(parent_task_id, parent_session_id)
        result = await self.scheduler.run(
            graph,
            parent_task_id=parent_task_id,
            parent_session_id=parent_session_id,
            parent_run_id=parent_run_id,
            provider_bundle=provider_bundle,
            provider_bundle_factory=provider_bundle_factory,
            inherited_pinned_skills=inherited_pinned_skills,
            inherited_pinned_skill_contexts=inherited_pinned_skill_contexts,
            allow_candidate_generation=allow_candidate_generation,
        )
        self._attach_runs_to_parent_task(result)
        return result

    def run(self, task: str) -> str:
        """Compatibility shim for old callers that only expected a string."""

        return f"team service requires run_team() for coordinated execution: {task}"

    def _validate_parent_task(self, parent_task_id: str | None, parent_session_id: str) -> None:
        if not parent_task_id:
            return
        loaded = self.loop.boot()
        task_service = getattr(loaded, "task_service", None)
        if task_service is None:
            raise RuntimeError("TeamService requires task_service when parent_task_id is provided")
        task = task_service.get_task(parent_task_id)
        if task is None:
            raise ValueError(f"Unknown parent_task_id: {parent_task_id}")
        if task.session_id != parent_session_id:
            raise ValueError(
                f"parent_task_id {parent_task_id!r} belongs to session {task.session_id!r}, "
                f"not {parent_session_id!r}"
            )

    def _attach_runs_to_parent_task(self, result: TeamRunResult) -> None:
        if not result.task_id or not result.run_ids:
            return
        loaded = self.loop.boot()
        task_service = getattr(loaded, "task_service", None)
        if task_service is None or task_service.get_task(result.task_id) is None:
            return
        run_store = getattr(loaded, "run_memory_store", None)
        for run_id in result.run_ids:
            skill_names: list[str] = []
            if run_store is not None:
                for record in run_store.list_runs():
                    if record.run_id == run_id:
                        skill_names = [receipt.skill_name for receipt in record.activated_skills]
                        break
            task_service.append_run(result.task_id, run_id, skill_names=skill_names)