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

from __future__ import annotations
from collections import defaultdict
from typing import Any, Callable, Self
from adk_fluent._base import BuilderBase
from google.adk.agents.loop_agent import LoopAgent
from google.adk.agents.parallel_agent import ParallelAgent
from google.adk.agents.sequential_agent import SequentialAgent

# ======================================================================
# Builder: Loop
# ======================================================================

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

    # --- Class-level alias / field maps ---
    _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)

    # --- Ergonomic aliases ---

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

    # --- Additive callback methods ---

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


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


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


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

    # --- Explicit field methods ---

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


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

    # --- Extra methods ---

    def step(self, agent: BaseAgent | AgentBuilder) -> Self:
        """Append an agent as the next step (lazy — built at .build() time)."""
        self._lists["sub_agents"].append(agent)
        return self

    # --- Dynamic field forwarding (safety net) ---

    def __getattr__(self, name: str):
        """Forward unknown methods to LoopAgent.model_fields for zero-maintenance compatibility."""
        if name.startswith("_"):
            raise AttributeError(name)

        # Resolve through alias map (class-level constants)
        _ALIASES = self.__class__._ALIASES
        _CALLBACK_ALIASES = self.__class__._CALLBACK_ALIASES
        _ADDITIVE_FIELDS = self.__class__._ADDITIVE_FIELDS

        field_name = _ALIASES.get(name, name)

        # Check if it's a callback alias
        if name in _CALLBACK_ALIASES:
            cb_field = _CALLBACK_ALIASES[name]
            def _cb_setter(fn: Callable) -> Self:
                self._callbacks[cb_field].append(fn)
                return self
            return _cb_setter

        # Validate against actual Pydantic schema
        if field_name not in LoopAgent.model_fields:
            available = sorted(
                set(LoopAgent.model_fields.keys())
                | set(_ALIASES.keys())
                | set(_CALLBACK_ALIASES.keys())
            )
            raise AttributeError(
                f"'{name}' is not a recognized field on LoopAgent. "
                f"Available: {', '.join(available)}"
            )

        # Return a setter that stores value and returns self for chaining
        def _setter(value: Any) -> Self:
            if field_name in _ADDITIVE_FIELDS:
                self._callbacks[field_name].append(value)
            else:
                self._config[field_name] = value
            return self

        return _setter

    # --- Terminal methods ---

    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()
        return LoopAgent(**config)


# ======================================================================
# Builder: FanOut
# ======================================================================

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

    # --- Class-level alias / field maps ---
    _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)

    # --- Ergonomic aliases ---

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

    # --- Additive callback methods ---

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


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


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


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

    # --- Explicit field methods ---

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

    # --- Extra methods ---

    def branch(self, agent: BaseAgent | AgentBuilder) -> Self:
        """Add a parallel branch agent (lazy — built at .build() time)."""
        self._lists["sub_agents"].append(agent)
        return self

    # --- Dynamic field forwarding (safety net) ---

    def __getattr__(self, name: str):
        """Forward unknown methods to ParallelAgent.model_fields for zero-maintenance compatibility."""
        if name.startswith("_"):
            raise AttributeError(name)

        # Resolve through alias map (class-level constants)
        _ALIASES = self.__class__._ALIASES
        _CALLBACK_ALIASES = self.__class__._CALLBACK_ALIASES
        _ADDITIVE_FIELDS = self.__class__._ADDITIVE_FIELDS

        field_name = _ALIASES.get(name, name)

        # Check if it's a callback alias
        if name in _CALLBACK_ALIASES:
            cb_field = _CALLBACK_ALIASES[name]
            def _cb_setter(fn: Callable) -> Self:
                self._callbacks[cb_field].append(fn)
                return self
            return _cb_setter

        # Validate against actual Pydantic schema
        if field_name not in ParallelAgent.model_fields:
            available = sorted(
                set(ParallelAgent.model_fields.keys())
                | set(_ALIASES.keys())
                | set(_CALLBACK_ALIASES.keys())
            )
            raise AttributeError(
                f"'{name}' is not a recognized field on ParallelAgent. "
                f"Available: {', '.join(available)}"
            )

        # Return a setter that stores value and returns self for chaining
        def _setter(value: Any) -> Self:
            if field_name in _ADDITIVE_FIELDS:
                self._callbacks[field_name].append(value)
            else:
                self._config[field_name] = value
            return self

        return _setter

    # --- Terminal methods ---

    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()
        return ParallelAgent(**config)


# ======================================================================
# Builder: Pipeline
# ======================================================================

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

    # --- Class-level alias / field maps ---
    _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)

    # --- Ergonomic aliases ---

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

    # --- Additive callback methods ---

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


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


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


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

    # --- Explicit field methods ---

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

    # --- Extra methods ---

    def step(self, agent: BaseAgent | AgentBuilder) -> Self:
        """Append an agent as the next step (lazy — built at .build() time)."""
        self._lists["sub_agents"].append(agent)
        return self

    # --- Dynamic field forwarding (safety net) ---

    def __getattr__(self, name: str):
        """Forward unknown methods to SequentialAgent.model_fields for zero-maintenance compatibility."""
        if name.startswith("_"):
            raise AttributeError(name)

        # Resolve through alias map (class-level constants)
        _ALIASES = self.__class__._ALIASES
        _CALLBACK_ALIASES = self.__class__._CALLBACK_ALIASES
        _ADDITIVE_FIELDS = self.__class__._ADDITIVE_FIELDS

        field_name = _ALIASES.get(name, name)

        # Check if it's a callback alias
        if name in _CALLBACK_ALIASES:
            cb_field = _CALLBACK_ALIASES[name]
            def _cb_setter(fn: Callable) -> Self:
                self._callbacks[cb_field].append(fn)
                return self
            return _cb_setter

        # Validate against actual Pydantic schema
        if field_name not in SequentialAgent.model_fields:
            available = sorted(
                set(SequentialAgent.model_fields.keys())
                | set(_ALIASES.keys())
                | set(_CALLBACK_ALIASES.keys())
            )
            raise AttributeError(
                f"'{name}' is not a recognized field on SequentialAgent. "
                f"Available: {', '.join(available)}"
            )

        # Return a setter that stores value and returns self for chaining
        def _setter(value: Any) -> Self:
            if field_name in _ADDITIVE_FIELDS:
                self._callbacks[field_name].append(value)
            else:
                self._config[field_name] = value
            return self

        return _setter

    # --- Terminal methods ---

    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()
        return SequentialAgent(**config)
