"""Auto-generated by adk-fluent generator. Manual edits will be overwritten."""

from __future__ import annotations

from collections import defaultdict
from collections.abc import Callable
from typing import TYPE_CHECKING, Any, Self

from google.adk.agents.loop_agent import LoopAgent
from google.adk.agents.parallel_agent import ParallelAgent
from google.adk.agents.sequential_agent import SequentialAgent

from adk_fluent._base import BuilderBase

if TYPE_CHECKING:
    from google.adk.agents.base_agent import BaseAgent


class Loop(BuilderBase):
    """A shell agent that run its sub-agents in a loop."""

    _ALIASES: dict[str, str] = {"describe": "description"}
    _CALLBACK_ALIASES: dict[str, str] = {"after_agent": "after_agent_callback", "before_agent": "before_agent_callback"}
    _ADDITIVE_FIELDS: set[str] = {"after_agent_callback", "before_agent_callback"}
    _ADK_TARGET_CLASS = LoopAgent

    def __init__(self, name: str) -> None:
        self._config: dict[str, Any] = {"name": name}
        self._callbacks: dict[str, list[Callable]] = defaultdict(list)
        self._lists: dict[str, list] = defaultdict(list)
        self._frozen = False

    def describe(self, value: str) -> Self:
        """Set the `description` field."""
        self = self._maybe_fork_for_mutation()
        self._config["description"] = value
        return self

    def after_agent(self, *fns: Callable[..., Any]) -> Self:
        """Append callback(s) to `after_agent_callback`. Multiple calls accumulate."""
        self = self._maybe_fork_for_mutation()
        for fn in fns:
            self._callbacks["after_agent_callback"].append(fn)
        return self

    def after_agent_if(self, condition: bool, fn: Callable[..., Any]) -> Self:
        """Append callback to `after_agent_callback` only if condition is True."""
        self = self._maybe_fork_for_mutation()
        if condition:
            self._callbacks["after_agent_callback"].append(fn)
        return self

    def before_agent(self, *fns: Callable[..., Any]) -> Self:
        """Append callback(s) to `before_agent_callback`. Multiple calls accumulate."""
        self = self._maybe_fork_for_mutation()
        for fn in fns:
            self._callbacks["before_agent_callback"].append(fn)
        return self

    def before_agent_if(self, condition: bool, fn: Callable[..., Any]) -> Self:
        """Append callback to `before_agent_callback` only if condition is True."""
        self = self._maybe_fork_for_mutation()
        if condition:
            self._callbacks["before_agent_callback"].append(fn)
        return self

    def sub_agents(self, value: list[BaseAgent]) -> Self:
        """Set the ``sub_agents`` field."""
        self = self._maybe_fork_for_mutation()
        self._config["sub_agents"] = value
        return self

    def max_iterations(self, value: int | None) -> Self:
        """Set the ``max_iterations`` field."""
        self = self._maybe_fork_for_mutation()
        self._config["max_iterations"] = value
        return self

    def step(self, value: BaseAgent) -> Self:
        """Append to ``sub_agents`` (lazy — built at .build() time)."""
        self = self._maybe_fork_for_mutation()
        self._lists["sub_agents"].append(value)
        return self

    def sub_agent(self, value: BaseAgent) -> Self:
        """Append to ``sub_agents`` (lazy — built at .build() time)."""
        self = self._maybe_fork_for_mutation()
        self._lists["sub_agents"].append(value)
        return self

    def to_ir(self) -> Any:
        """Convert this Loop builder to a LoopNode IR node."""
        from adk_fluent._helpers import _loop_to_ir

        return _loop_to_ir(self)

    def build(self) -> LoopAgent:
        """A shell agent that run its sub-agents in a loop. Resolve into a native ADK LoopAgent."""
        config = self._prepare_build_config()
        result = self._safe_build(LoopAgent, config)
        return self._apply_native_hooks(result)


class FanOut(BuilderBase):
    """A shell agent that runs its sub-agents in parallel in an isolated manner."""

    _ALIASES: dict[str, str] = {"describe": "description"}
    _CALLBACK_ALIASES: dict[str, str] = {"after_agent": "after_agent_callback", "before_agent": "before_agent_callback"}
    _ADDITIVE_FIELDS: set[str] = {"after_agent_callback", "before_agent_callback"}
    _ADK_TARGET_CLASS = ParallelAgent

    def __init__(self, name: str) -> None:
        self._config: dict[str, Any] = {"name": name}
        self._callbacks: dict[str, list[Callable]] = defaultdict(list)
        self._lists: dict[str, list] = defaultdict(list)
        self._frozen = False

    def describe(self, value: str) -> Self:
        """Set the `description` field."""
        self = self._maybe_fork_for_mutation()
        self._config["description"] = value
        return self

    def after_agent(self, *fns: Callable[..., Any]) -> Self:
        """Append callback(s) to `after_agent_callback`. Multiple calls accumulate."""
        self = self._maybe_fork_for_mutation()
        for fn in fns:
            self._callbacks["after_agent_callback"].append(fn)
        return self

    def after_agent_if(self, condition: bool, fn: Callable[..., Any]) -> Self:
        """Append callback to `after_agent_callback` only if condition is True."""
        self = self._maybe_fork_for_mutation()
        if condition:
            self._callbacks["after_agent_callback"].append(fn)
        return self

    def before_agent(self, *fns: Callable[..., Any]) -> Self:
        """Append callback(s) to `before_agent_callback`. Multiple calls accumulate."""
        self = self._maybe_fork_for_mutation()
        for fn in fns:
            self._callbacks["before_agent_callback"].append(fn)
        return self

    def before_agent_if(self, condition: bool, fn: Callable[..., Any]) -> Self:
        """Append callback to `before_agent_callback` only if condition is True."""
        self = self._maybe_fork_for_mutation()
        if condition:
            self._callbacks["before_agent_callback"].append(fn)
        return self

    def sub_agents(self, value: list[BaseAgent]) -> Self:
        """Set the ``sub_agents`` field."""
        self = self._maybe_fork_for_mutation()
        self._config["sub_agents"] = value
        return self

    def branch(self, value: BaseAgent) -> Self:
        """Append to ``sub_agents`` (lazy — built at .build() time)."""
        self = self._maybe_fork_for_mutation()
        self._lists["sub_agents"].append(value)
        return self

    def sub_agent(self, value: BaseAgent) -> Self:
        """Append to ``sub_agents`` (lazy — built at .build() time)."""
        self = self._maybe_fork_for_mutation()
        self._lists["sub_agents"].append(value)
        return self

    def to_ir(self) -> Any:
        """Convert this FanOut builder to a ParallelNode IR node."""
        from adk_fluent._helpers import _fanout_to_ir

        return _fanout_to_ir(self)

    def build(self) -> ParallelAgent:
        """A shell agent that runs its sub-agents in parallel in an isolated manner. Resolve into a native ADK ParallelAgent."""
        config = self._prepare_build_config()
        result = self._safe_build(ParallelAgent, config)
        return self._apply_native_hooks(result)


class Pipeline(BuilderBase):
    """A shell agent that runs its sub-agents in sequence."""

    _ALIASES: dict[str, str] = {"describe": "description"}
    _CALLBACK_ALIASES: dict[str, str] = {"after_agent": "after_agent_callback", "before_agent": "before_agent_callback"}
    _ADDITIVE_FIELDS: set[str] = {"after_agent_callback", "before_agent_callback"}
    _ADK_TARGET_CLASS = SequentialAgent

    def __init__(self, name: str) -> None:
        self._config: dict[str, Any] = {"name": name}
        self._callbacks: dict[str, list[Callable]] = defaultdict(list)
        self._lists: dict[str, list] = defaultdict(list)
        self._frozen = False

    def describe(self, value: str) -> Self:
        """Set the `description` field."""
        self = self._maybe_fork_for_mutation()
        self._config["description"] = value
        return self

    def after_agent(self, *fns: Callable[..., Any]) -> Self:
        """Append callback(s) to `after_agent_callback`. Multiple calls accumulate."""
        self = self._maybe_fork_for_mutation()
        for fn in fns:
            self._callbacks["after_agent_callback"].append(fn)
        return self

    def after_agent_if(self, condition: bool, fn: Callable[..., Any]) -> Self:
        """Append callback to `after_agent_callback` only if condition is True."""
        self = self._maybe_fork_for_mutation()
        if condition:
            self._callbacks["after_agent_callback"].append(fn)
        return self

    def before_agent(self, *fns: Callable[..., Any]) -> Self:
        """Append callback(s) to `before_agent_callback`. Multiple calls accumulate."""
        self = self._maybe_fork_for_mutation()
        for fn in fns:
            self._callbacks["before_agent_callback"].append(fn)
        return self

    def before_agent_if(self, condition: bool, fn: Callable[..., Any]) -> Self:
        """Append callback to `before_agent_callback` only if condition is True."""
        self = self._maybe_fork_for_mutation()
        if condition:
            self._callbacks["before_agent_callback"].append(fn)
        return self

    def sub_agents(self, value: list[BaseAgent]) -> Self:
        """Set the ``sub_agents`` field."""
        self = self._maybe_fork_for_mutation()
        self._config["sub_agents"] = value
        return self

    def step(self, value: BaseAgent) -> Self:
        """Append to ``sub_agents`` (lazy — built at .build() time)."""
        self = self._maybe_fork_for_mutation()
        self._lists["sub_agents"].append(value)
        return self

    def sub_agent(self, value: BaseAgent) -> Self:
        """Append to ``sub_agents`` (lazy — built at .build() time)."""
        self = self._maybe_fork_for_mutation()
        self._lists["sub_agents"].append(value)
        return self

    def to_ir(self) -> Any:
        """Convert this Pipeline builder to a SequenceNode IR node."""
        from adk_fluent._helpers import _pipeline_to_ir

        return _pipeline_to_ir(self)

    def build(self) -> SequentialAgent:
        """A shell agent that runs its sub-agents in sequence. Resolve into a native ADK SequentialAgent."""
        config = self._prepare_build_config()
        result = self._safe_build(SequentialAgent, config)
        return self._apply_native_hooks(result)
