"""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.base_agent import BaseAgent as _ADK_BaseAgent
from google.adk.agents.llm_agent import LlmAgent

from adk_fluent._base import BuilderBase

if TYPE_CHECKING:
    from collections.abc import AsyncIterator, Awaitable
    from typing import Literal

    from google.adk.agents.readonly_context import ReadonlyContext
    from google.adk.code_executors.base_code_executor import BaseCodeExecutor
    from google.adk.models.base_llm import BaseLlm
    from google.adk.planners.base_planner import BasePlanner
    from google.adk.tools.base_tool import BaseTool
    from google.adk.tools.base_toolset import BaseToolset
    from google.genai.types import Content, File, GenerateContentConfig, Part
    from pydantic import BaseModel


class BaseAgent(BuilderBase):
    """Base class for all agents in Agent Development Kit."""

    _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 = _ADK_BaseAgent

    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 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 build(self) -> _ADK_BaseAgent:
        """Base class for all agents in Agent Development Kit. Resolve into a native ADK _ADK_BaseAgent."""
        config = self._prepare_build_config()
        result = self._safe_build(_ADK_BaseAgent, config)
        return self._apply_native_hooks(result)


class Agent(BuilderBase):
    """LLM-based Agent."""

    _ALIASES: dict[str, str] = {
        "describe": "description",
        "global_instruct": "global_instruction",
        "instruct": "instruction",
        "save_as": "output_key",
        "static": "static_instruction",
    }
    _CALLBACK_ALIASES: dict[str, str] = {
        "after_agent": "after_agent_callback",
        "after_model": "after_model_callback",
        "after_tool": "after_tool_callback",
        "before_agent": "before_agent_callback",
        "before_model": "before_model_callback",
        "before_tool": "before_tool_callback",
        "on_model_error": "on_model_error_callback",
        "on_tool_error": "on_tool_error_callback",
    }
    _ADDITIVE_FIELDS: set[str] = {
        "after_agent_callback",
        "after_model_callback",
        "after_tool_callback",
        "before_agent_callback",
        "before_model_callback",
        "before_tool_callback",
        "on_model_error_callback",
        "on_tool_error_callback",
    }
    _ADK_TARGET_CLASS = LlmAgent

    def __init__(self, name: str, model: str | None = None) -> 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
        if model is not None:
            self._config["model"] = model

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

    def global_instruct(self, value: str | Callable[[ReadonlyContext], str | Awaitable[str]]) -> Self:
        """Set the `global_instruction` field."""
        self = self._maybe_fork_for_mutation()
        self._config["global_instruction"] = value
        return self

    def instruct(self, value: str | Callable[[ReadonlyContext], str | Awaitable[str]]) -> Self:
        """Set the `instruction` field."""
        self = self._maybe_fork_for_mutation()
        self._config["instruction"] = value
        return self

    def save_as(self, value: str | None) -> Self:
        """Session state key where the agent's response text is stored. Downstream agents and state transforms can read this key via ``state['key']``. Alias for ``output_key``. Prefer ``.writes(key)``."""
        self = self._maybe_fork_for_mutation()
        self._config["output_key"] = value
        return self

    def static(self, value: Content | str | File | Part | list[str | File | Part] | None) -> Self:
        """Set the `static_instruction` field."""
        self = self._maybe_fork_for_mutation()
        self._config["static_instruction"] = value
        return self

    def history(self, value: Literal["default", "none"]) -> Self:
        """Deprecated: use ``.context()`` instead."""
        self = self._maybe_fork_for_mutation()
        import warnings

        warnings.warn(
            ".history() is deprecated, use .context() instead",
            DeprecationWarning,
            stacklevel=2,
        )
        self._config["include_contents"] = value
        return self

    def include_history(self, value: Literal["default", "none"]) -> Self:
        """Deprecated: use ``.context()`` instead."""
        self = self._maybe_fork_for_mutation()
        import warnings

        warnings.warn(
            ".include_history() is deprecated, use .context() instead",
            DeprecationWarning,
            stacklevel=2,
        )
        self._config["include_contents"] = value
        return self

    def outputs(self, value: str | None) -> Self:
        """Deprecated: use ``.save_as()`` instead."""
        self = self._maybe_fork_for_mutation()
        import warnings

        warnings.warn(
            ".outputs() is deprecated, use .save_as() instead",
            DeprecationWarning,
            stacklevel=2,
        )
        self._config["output_key"] = value
        return self

    def static_instruct(self, value: Content | str | File | Part | list[str | File | Part] | None) -> Self:
        """Deprecated: use ``.static()`` instead."""
        self = self._maybe_fork_for_mutation()
        import warnings

        warnings.warn(
            ".static_instruct() is deprecated, use .static() instead",
            DeprecationWarning,
            stacklevel=2,
        )
        self._config["static_instruction"] = 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 after_model(self, *fns: Callable[..., Any]) -> Self:
        """Append callback(s) to `after_model_callback`. Multiple calls accumulate."""
        self = self._maybe_fork_for_mutation()
        for fn in fns:
            self._callbacks["after_model_callback"].append(fn)
        return self

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

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

    def after_tool_if(self, condition: bool, fn: Callable[..., Any]) -> Self:
        """Append callback to `after_tool_callback` only if condition is True."""
        self = self._maybe_fork_for_mutation()
        if condition:
            self._callbacks["after_tool_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 before_model(self, *fns: Callable[..., Any]) -> Self:
        """Append callback(s) to `before_model_callback`. Multiple calls accumulate."""
        self = self._maybe_fork_for_mutation()
        for fn in fns:
            self._callbacks["before_model_callback"].append(fn)
        return self

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

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

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

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

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

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

    def on_tool_error_if(self, condition: bool, fn: Callable[..., Any]) -> Self:
        """Append callback to `on_tool_error_callback` only if condition is True."""
        self = self._maybe_fork_for_mutation()
        if condition:
            self._callbacks["on_tool_error_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 model(self, value: str | BaseLlm) -> Self:
        """Set the ``model`` field."""
        self = self._maybe_fork_for_mutation()
        self._config["model"] = value
        return self

    def tools(self, value: list[Callable[..., Any] | BaseTool | BaseToolset]) -> Self:
        """Set the ``tools`` field."""
        self = self._maybe_fork_for_mutation()
        self._config["tools"] = value
        return self

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

    def disallow_transfer_to_parent(self, value: bool) -> Self:
        """Prevent this agent from transferring control back to its parent. Also forces a handoff back to parent on the next turn, preventing the user from getting stuck. See also ``.isolate()``."""
        self = self._maybe_fork_for_mutation()
        self._config["disallow_transfer_to_parent"] = value
        return self

    def disallow_transfer_to_peers(self, value: bool) -> Self:
        """Prevent this agent from transferring control to sibling agents. See also ``.isolate()``."""
        self = self._maybe_fork_for_mutation()
        self._config["disallow_transfer_to_peers"] = value
        return self

    def input_schema(self, value: type[BaseModel] | None) -> Self:
        """Schema defining the expected input structure when this agent is invoked as a tool by another agent.

        When another agent invokes this agent via ``AgentTool``, the calling
        agent's arguments are validated against this Pydantic model. Irrelevant
        for top-level agents — only for agents that serve as tools.

        .. note::

           Prefer ``.accepts(Model)`` over this method for clarity.
           ``.accepts()`` is the recommended alias on BuilderBase.

           - ``.accepts(Model)`` → tool-mode input validation (same as this)
           - ``.consumes(Model)`` → contract annotation (no runtime effect)
        """
        self = self._maybe_fork_for_mutation()
        self._config["input_schema"] = value
        return self

    def output_schema(self, value: type[BaseModel] | None) -> Self:
        """Force the LLM to respond with structured JSON matching a Pydantic model.

        When set, the agent replies **only** with JSON data conforming to
        this schema. The agent **cannot use tools** while ``output_schema``
        is active.

        .. note::

           Prefer ``.returns(Model)`` or ``@ Model`` over this method.
           ``.returns()`` sets the same ADK constraint AND automatically
           parses the response in ``.ask()`` calls. This method sets the
           raw ADK field without automatic parsing.

           - ``.returns(Model)`` / ``@ Model`` → LLM constraint + parsing
           - ``.output_schema(Model)`` → LLM constraint only (raw field)
           - ``.writes(key)`` → stores text in state (no format constraint)
           - ``.produces(Model)`` → contract annotation (no runtime effect)
        """
        self = self._maybe_fork_for_mutation()
        self._config["output_schema"] = value
        return self

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

    def code_executor(self, value: BaseCodeExecutor | None) -> Self:
        """Set the ``code_executor`` field."""
        self = self._maybe_fork_for_mutation()
        self._config["code_executor"] = 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 tool(self, fn_or_tool: Any, *, require_confirmation: bool = False) -> Self:
        """Add a single tool (appends). Wraps plain callables in FunctionTool when require_confirmation=True."""
        from adk_fluent._helpers import _add_tool

        return _add_tool(self, fn_or_tool, require_confirmation=require_confirmation)

    def guardrail(self, fn: Callable[..., Any]) -> Self:
        """Attach a guardrail function as both before_model and after_model callback."""
        self = self._maybe_fork_for_mutation()
        self._callbacks["before_model_callback"].append(fn)
        self._callbacks["after_model_callback"].append(fn)
        return self

    def ask(self, prompt: str) -> str:
        """One-shot execution. Build agent, send prompt, return response text."""
        from adk_fluent._helpers import run_one_shot

        return run_one_shot(self, prompt)

    async def ask_async(self, prompt: str) -> str:
        """Async one-shot execution."""
        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]:
        """Streaming execution. Yields response text chunks."""
        from adk_fluent._helpers import run_stream

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

    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 session(self) -> Any:
        """Create an interactive session context manager. Use with 'async with'."""
        from adk_fluent._helpers import create_session

        return create_session(self)

    def map(self, prompts: list[str], *, concurrency: int = 5) -> list[str]:
        """Run agent 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]:
        """Async batch execution against multiple prompts."""
        from adk_fluent._helpers import run_map_async

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

    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 context(self, spec: Any) -> Self:
        """Declare what conversation context this agent should see. Accepts a C module transform (C.none(), C.user_only(), C.from_state(), etc.)."""
        self = self._maybe_fork_for_mutation()
        self._config["_context_spec"] = spec
        return self

    def show(self) -> Self:
        """Force this agent's events to be user-facing (override topology inference)."""
        from adk_fluent._helpers import _show_agent

        return _show_agent(self)

    def hide(self) -> Self:
        """Force this agent's events to be internal (override topology inference)."""
        from adk_fluent._helpers import _hide_agent

        return _hide_agent(self)

    def memory(self, mode: str = "preload") -> Self:
        """Add memory tools to this agent. Modes: 'preload', 'on_demand', 'both'."""
        from adk_fluent._helpers import _add_memory

        return _add_memory(self, mode)

    def memory_auto_save(self) -> Self:
        """Auto-save session to memory after each agent run."""
        from adk_fluent._helpers import _add_memory_auto_save

        return _add_memory_auto_save(self)

    def delegate(self, agent: Any) -> Self:
        """Add an agent as a delegatable tool (wraps in AgentTool). The coordinator LLM can route to this agent."""
        from adk_fluent._helpers import delegate_agent

        return delegate_agent(self, agent)

    def isolate(self) -> Self:
        """Prevent this agent from transferring to parent or peers. Sets both disallow_transfer_to_parent and disallow_transfer_to_peers to True. Use for specialist agents that should complete their task and return."""
        from adk_fluent._helpers import _isolate_agent

        return _isolate_agent(
            self,
        )

    def stay(self) -> Self:
        """Prevent this agent from transferring back to its parent. Use for agents that should complete their work before returning."""
        from adk_fluent._helpers import _stay_agent

        return _stay_agent(
            self,
        )

    def no_peers(self) -> Self:
        """Prevent this agent from transferring to sibling agents. The agent can still return to its parent."""
        from adk_fluent._helpers import _no_peers_agent

        return _no_peers_agent(
            self,
        )

    def tool_schema(self, schema: type) -> Self:
        """Attach a ToolSchema declaring tool state dependencies."""
        self = self._maybe_fork_for_mutation()
        self._config["_tool_schema"] = schema
        return self

    def callback_schema(self, schema: type) -> Self:
        """Attach a CallbackSchema declaring callback state dependencies."""
        self = self._maybe_fork_for_mutation()
        self._config["_callback_schema"] = schema
        return self

    def prompt_schema(self, schema: type) -> Self:
        """Attach a PromptSchema declaring prompt state dependencies."""
        self = self._maybe_fork_for_mutation()
        self._config["_prompt_schema"] = schema
        return self

    def to_ir(self) -> Any:
        """Convert this Agent builder to an AgentNode IR node."""
        from adk_fluent._helpers import _agent_to_ir

        return _agent_to_ir(self)

    def build(self) -> LlmAgent:
        """LLM-based Agent. Resolve into a native ADK LlmAgent."""
        config = self._prepare_build_config()
        result = self._safe_build(LlmAgent, config)
        return self._apply_native_hooks(result)
