"""Default preset orchestrator builders."""

from __future__ import annotations

from typing import Any, Literal

from agentbyte.llm import BaseChatCompletionClient
from agentbyte.orchestration import (
    AIOrchestrator,
    AgentHistoryPolicy,
    FinalResultPolicy,
    HistoryContextPolicy,
    PlanBasedOrchestrator,
    ReviewGatePolicy,
    RoundRobinOrchestrator,
    SelectorContextPolicy,
)
from agentbyte.termination import (
    BaseTermination,
    CompositeTermination,
    MaxMessageTermination,
    TextMentionTermination,
)

from .agents import get_researcher, get_reviewer, get_writer
from .clients import build_chat_client


def _resolve_model_client(
    model_client: BaseChatCompletionClient | None = None,
    *,
    provider: str | None = None,
    model: str | None = None,
    config: dict[str, Any] | None = None,
) -> BaseChatCompletionClient:
    if model_client is not None:
        return model_client
    return build_chat_client(provider=provider, model=model, config=config)


def get_round_robin_orchestrator(
    model_client: BaseChatCompletionClient | None = None,
    *,
    provider: str | None = None,
    model: str | None = None,
    config: dict[str, Any] | None = None,
    termination: BaseTermination | None = None,
    max_iterations: int = 6,
    name: str = "round_robin_research_writer",
    description: str = "Default round-robin researcher/writer collaboration preset.",
    **kwargs: Any,
) -> RoundRobinOrchestrator:
    """Build the default round-robin orchestrator preset."""
    del kwargs
    client = _resolve_model_client(
        model_client,
        provider=provider,
        model=model,
        config=config,
    )
    orchestrator = RoundRobinOrchestrator(
        agents=[
            get_researcher(client),
            get_writer(client, completion_signal="TERMINATE"),
        ],
        termination=termination
        or CompositeTermination(
            [TextMentionTermination("TERMINATE"), MaxMessageTermination(6)],
            mode="any",
        ),
        max_iterations=max_iterations,
        history_policy=HistoryContextPolicy(
            default=AgentHistoryPolicy(window="full", format="text")
        ),
        final_result_policy=FinalResultPolicy(
            prefer_agent="writer",
            strip_signals=["TERMINATE"],
            strip_mode="trailing",
        ),
    )
    orchestrator.name = name
    orchestrator.description = description
    orchestrator.example_tasks = [
        "Write a brief article about solar energy benefits.",
        "Create a short note on renewable energy advantages.",
    ]
    return orchestrator


def get_ai_orchestrator(
    model_client: BaseChatCompletionClient | None = None,
    *,
    provider: str | None = None,
    model: str | None = None,
    config: dict[str, Any] | None = None,
    termination: BaseTermination | None = None,
    max_iterations: int = 8,
    name: str = "ai_research_writer_reviewer",
    description: str = "Default AI-driven researcher/writer/reviewer collaboration preset.",
    **kwargs: Any,
) -> AIOrchestrator:
    """Build the default AI-driven orchestrator preset."""
    del kwargs
    client = _resolve_model_client(
        model_client,
        provider=provider,
        model=model,
        config=config,
    )
    orchestrator = AIOrchestrator(
        agents=[
            get_researcher(client),
            get_writer(client),
            get_reviewer(client, approval_signal="APPROVED"),
        ],
        termination=termination
        or CompositeTermination(
            [
                TextMentionTermination("APPROVED"),
                TextMentionTermination("TERMINATE"),
                MaxMessageTermination(8),
            ],
            mode="any",
        ),
        model_client=client,
        max_iterations=max_iterations,
        history_policy=HistoryContextPolicy(
            default=AgentHistoryPolicy(window="full", format="text")
        ),
        selector_policy=SelectorContextPolicy(
            window="recent",
            message_limit=12,
            format="summary",
        ),
        final_result_policy=FinalResultPolicy(
            prefer_agent="writer",
            strip_signals=["APPROVED", "TERMINATE"],
            strip_mode="trailing",
        ),
    )
    orchestrator.name = name
    orchestrator.description = description
    orchestrator.example_tasks = [
        "Write a brief article about solar energy benefits.",
        "Draft a concise note on remote work benefits.",
    ]
    return orchestrator


def get_plan_based_orchestrator(
    model_client: BaseChatCompletionClient | None = None,
    *,
    provider: str | None = None,
    model: str | None = None,
    config: dict[str, Any] | None = None,
    termination: BaseTermination | None = None,
    max_iterations: int = 8,
    max_step_retries: int = 2,
    name: str = "plan_based_research_writer_reviewer",
    description: str = "Default plan-based researcher/writer/reviewer collaboration preset.",
    **kwargs: Any,
) -> PlanBasedOrchestrator:
    """Build the default plan-based orchestrator preset."""
    del kwargs
    client = _resolve_model_client(
        model_client,
        provider=provider,
        model=model,
        config=config,
    )
    orchestrator = PlanBasedOrchestrator(
        agents=[
            get_researcher(client),
            get_writer(client),
            get_reviewer(client, approval_signal="APPROVED"),
        ],
        termination=termination or MaxMessageTermination(8),
        model_client=client,
        max_iterations=max_iterations,
        max_step_retries=max_step_retries,
        history_policy=HistoryContextPolicy(
            default=AgentHistoryPolicy(
                window="recent",
                message_limit=5,
                format="messages",
            )
        ),
        review_gate_policy=ReviewGatePolicy(
            enabled=True,
            reviewer_agent_name="reviewer",
            revision_agent_name="writer",
            approval_signal="APPROVED",
            max_revision_rounds=max_step_retries,
        ),
        final_result_policy=FinalResultPolicy(
            prefer_agent="writer",
            strip_signals=["APPROVED"],
            strip_mode="trailing",
        ),
    )
    orchestrator.name = name
    orchestrator.description = description
    orchestrator.example_tasks = [
        "Write a brief article about solar energy benefits.",
        "Draft a concise note on remote work benefits.",
    ]
    return orchestrator


def get_orchestrator(
    pattern: Literal["round_robin", "ai", "plan_based"] = "ai",
    model_client: BaseChatCompletionClient | None = None,
    **kwargs: Any,
) -> RoundRobinOrchestrator | AIOrchestrator | PlanBasedOrchestrator:
    """Build one of the default orchestrator presets."""
    if pattern == "round_robin":
        return get_round_robin_orchestrator(model_client=model_client, **kwargs)
    if pattern == "ai":
        return get_ai_orchestrator(model_client=model_client, **kwargs)
    if pattern == "plan_based":
        return get_plan_based_orchestrator(model_client=model_client, **kwargs)
    raise ValueError(
        "Unsupported orchestrator pattern. Expected 'round_robin', 'ai', or 'plan_based'."
    )


__all__ = [
    "get_ai_orchestrator",
    "get_orchestrator",
    "get_plan_based_orchestrator",
    "get_round_robin_orchestrator",
]
