"""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 adk_fluent._base import BuilderBase

if TYPE_CHECKING:
    from collections.abc import AsyncIterator

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


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"}

    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 agent description (metadata for transfer routing and topology display — NOT sent to the LLM as instruction). Always set this on sub-agents so the coordinator LLM can pick the right specialist."""
        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 ask(self, prompt: str) -> str:
        """One-shot SYNC execution (blocking). Builds loop, sends prompt, returns response text."""
        from adk_fluent._helpers import run_one_shot

        return run_one_shot(self, prompt)

    async def ask_async(self, prompt: str) -> str:
        """One-shot ASYNC execution (non-blocking, use with await)."""
        from adk_fluent._helpers import run_one_shot_async

        return await run_one_shot_async(self, prompt)

    async def stream(self, prompt: str) -> AsyncIterator[str]:
        """ASYNC streaming execution. Yields response text chunks as they arrive."""
        from adk_fluent._helpers import run_stream

        async for chunk in run_stream(self, prompt):
            yield chunk

    async def events(self, prompt: str) -> AsyncIterator[Any]:
        """Stream raw ADK Event objects. Yields every event including state deltas and function calls."""
        from adk_fluent._helpers import run_events

        async for chunk in run_events(self, prompt):
            yield chunk

    def session(self) -> Any:
        """Create an interactive multi-turn chat session. Returns an async context manager."""
        from adk_fluent._helpers import create_session

        return create_session(self)

    def map(self, prompts: list[str], *, concurrency: int = 5) -> list[str]:
        """Batch SYNC execution (blocking). Run loop against multiple prompts with bounded concurrency."""
        from adk_fluent._helpers import run_map

        return run_map(self, prompts, concurrency=concurrency)

    async def map_async(self, prompts: list[str], *, concurrency: int = 5) -> list[str]:
        """Batch ASYNC execution (non-blocking, use with await)."""
        from adk_fluent._helpers import run_map_async

        return await run_map_async(self, prompts, concurrency=concurrency)

    def test(
        self, prompt: str, *, contains: str | None = None, matches: str | None = None, equals: str | None = None
    ) -> Self:
        """Run a smoke test. Calls .ask() internally, asserts output matches condition."""
        from adk_fluent._helpers import run_inline_test

        return run_inline_test(self, prompt, contains=contains, matches=matches, equals=equals)

    def eval(self, prompt: str, *, expect: str | None = None, criteria: Any | None = None) -> Any:
        """Inline evaluation. Run a single eval case against this loop."""
        from adk_fluent._helpers import _eval_inline

        return _eval_inline(self, prompt, expect=expect, criteria=criteria)

    def eval_suite(self) -> Any:
        """Create an evaluation suite builder for this loop."""
        from adk_fluent._helpers import _eval_suite

        return _eval_suite(self)

    def build(self) -> LoopAgent:
        """A shell agent that run its sub-agents in a loop. Resolve into a native ADK LoopAgent."""
        from google.adk.agents.loop_agent import 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"}

    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 agent description (metadata for transfer routing and topology display — NOT sent to the LLM as instruction). Always set this on sub-agents so the coordinator LLM can pick the right specialist."""
        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 ask(self, prompt: str) -> str:
        """One-shot SYNC execution (blocking). Builds fan-out, sends prompt, returns response text."""
        from adk_fluent._helpers import run_one_shot

        return run_one_shot(self, prompt)

    async def ask_async(self, prompt: str) -> str:
        """One-shot ASYNC execution (non-blocking, use with await)."""
        from adk_fluent._helpers import run_one_shot_async

        return await run_one_shot_async(self, prompt)

    async def stream(self, prompt: str) -> AsyncIterator[str]:
        """ASYNC streaming execution. Yields response text chunks as they arrive."""
        from adk_fluent._helpers import run_stream

        async for chunk in run_stream(self, prompt):
            yield chunk

    async def events(self, prompt: str) -> AsyncIterator[Any]:
        """Stream raw ADK Event objects. Yields every event including state deltas and function calls."""
        from adk_fluent._helpers import run_events

        async for chunk in run_events(self, prompt):
            yield chunk

    def session(self) -> Any:
        """Create an interactive multi-turn chat session. Returns an async context manager."""
        from adk_fluent._helpers import create_session

        return create_session(self)

    def map(self, prompts: list[str], *, concurrency: int = 5) -> list[str]:
        """Batch SYNC execution (blocking). Run fan-out against multiple prompts with bounded concurrency."""
        from adk_fluent._helpers import run_map

        return run_map(self, prompts, concurrency=concurrency)

    async def map_async(self, prompts: list[str], *, concurrency: int = 5) -> list[str]:
        """Batch ASYNC execution (non-blocking, use with await)."""
        from adk_fluent._helpers import run_map_async

        return await run_map_async(self, prompts, concurrency=concurrency)

    def test(
        self, prompt: str, *, contains: str | None = None, matches: str | None = None, equals: str | None = None
    ) -> Self:
        """Run a smoke test. Calls .ask() internally, asserts output matches condition."""
        from adk_fluent._helpers import run_inline_test

        return run_inline_test(self, prompt, contains=contains, matches=matches, equals=equals)

    def eval(self, prompt: str, *, expect: str | None = None, criteria: Any | None = None) -> Any:
        """Inline evaluation. Run a single eval case against this fan-out."""
        from adk_fluent._helpers import _eval_inline

        return _eval_inline(self, prompt, expect=expect, criteria=criteria)

    def eval_suite(self) -> Any:
        """Create an evaluation suite builder for this fan-out."""
        from adk_fluent._helpers import _eval_suite

        return _eval_suite(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."""
        from google.adk.agents.parallel_agent import 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"}

    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 agent description (metadata for transfer routing and topology display — NOT sent to the LLM as instruction). Always set this on sub-agents so the coordinator LLM can pick the right specialist."""
        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 ask(self, prompt: str) -> str:
        """One-shot SYNC execution (blocking). Builds pipeline, sends prompt, returns response text."""
        from adk_fluent._helpers import run_one_shot

        return run_one_shot(self, prompt)

    async def ask_async(self, prompt: str) -> str:
        """One-shot ASYNC execution (non-blocking, use with await)."""
        from adk_fluent._helpers import run_one_shot_async

        return await run_one_shot_async(self, prompt)

    async def stream(self, prompt: str) -> AsyncIterator[str]:
        """ASYNC streaming execution. Yields response text chunks as they arrive."""
        from adk_fluent._helpers import run_stream

        async for chunk in run_stream(self, prompt):
            yield chunk

    async def events(self, prompt: str) -> AsyncIterator[Any]:
        """Stream raw ADK Event objects. Yields every event including state deltas and function calls."""
        from adk_fluent._helpers import run_events

        async for chunk in run_events(self, prompt):
            yield chunk

    def session(self) -> Any:
        """Create an interactive multi-turn chat session. Returns an async context manager."""
        from adk_fluent._helpers import create_session

        return create_session(self)

    def map(self, prompts: list[str], *, concurrency: int = 5) -> list[str]:
        """Batch SYNC execution (blocking). Run pipeline against multiple prompts with bounded concurrency."""
        from adk_fluent._helpers import run_map

        return run_map(self, prompts, concurrency=concurrency)

    async def map_async(self, prompts: list[str], *, concurrency: int = 5) -> list[str]:
        """Batch ASYNC execution (non-blocking, use with await)."""
        from adk_fluent._helpers import run_map_async

        return await run_map_async(self, prompts, concurrency=concurrency)

    def test(
        self, prompt: str, *, contains: str | None = None, matches: str | None = None, equals: str | None = None
    ) -> Self:
        """Run a smoke test. Calls .ask() internally, asserts output matches condition."""
        from adk_fluent._helpers import run_inline_test

        return run_inline_test(self, prompt, contains=contains, matches=matches, equals=equals)

    def eval(self, prompt: str, *, expect: str | None = None, criteria: Any | None = None) -> Any:
        """Inline evaluation. Run a single eval case against this pipeline."""
        from adk_fluent._helpers import _eval_inline

        return _eval_inline(self, prompt, expect=expect, criteria=criteria)

    def eval_suite(self) -> Any:
        """Create an evaluation suite builder for this pipeline."""
        from adk_fluent._helpers import _eval_suite

        return _eval_suite(self)

    def build(self) -> SequentialAgent:
        """A shell agent that runs its sub-agents in sequence. Resolve into a native ADK SequentialAgent."""
        from google.adk.agents.sequential_agent import SequentialAgent

        config = self._prepare_build_config()
        result = self._safe_build(SequentialAgent, config)
        return self._apply_native_hooks(result)
