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

# ======================================================================
# Builder: BaseAgent
# ======================================================================

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

    # --- 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 ---

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

    def __getattr__(self, name: str):
        """Forward unknown methods to _ADK_BaseAgent.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 _ADK_BaseAgent.model_fields:
            available = sorted(
                set(_ADK_BaseAgent.model_fields.keys())
                | set(_ALIASES.keys())
                | set(_CALLBACK_ALIASES.keys())
            )
            raise AttributeError(
                f"'{name}' is not a recognized field on _ADK_BaseAgent. "
                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) -> _ADK_BaseAgent:
        """Base class for all agents in Agent Development Kit. Resolve into a native ADK _ADK_BaseAgent."""
        config = self._prepare_build_config()
        return _ADK_BaseAgent(**config)


# ======================================================================
# Builder: Agent
# ======================================================================

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

    # --- Class-level alias / field maps ---
    _ALIASES: dict[str, str] = {'describe': 'description', 'global_instruct': 'global_instruction', 'history': 'include_contents', 'instruct': 'instruction', 'outputs': '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_tool_callback', 'after_agent_callback', 'after_model_callback', 'before_agent_callback', 'before_tool_callback', 'before_model_callback', 'on_model_error_callback', 'on_tool_error_callback'}


    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)
        if model is not None:
            self._config["model"] = model

    # --- Ergonomic aliases ---

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


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


    def history(self, value: Literal[default, none]) -> Self:
        """Set the `include_contents` field."""
        self._config["include_contents"] = value
        return self


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


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


    def static(self, value: Union[Content, str, File, Part, list[Union[str, File, Part]], NoneType]) -> Self:
        """Set the `static_instruction` field."""
        self._config["static_instruction"] = 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 after_model(self, *fns: Callable) -> Self:
        """Append callback(s) to `after_model_callback`. Multiple calls accumulate."""
        for fn in fns:
            self._callbacks["after_model_callback"].append(fn)
        return self


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


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


    def after_tool_if(self, condition: bool, fn: Callable) -> Self:
        """Append callback to `after_tool_callback` only if condition is True."""
        if condition:
            self._callbacks["after_tool_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


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


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


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


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


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


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


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


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


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


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


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


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


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


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


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


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

    # --- Extra methods ---

    def tool(self, fn_or_tool: Callable | BaseTool) -> Self:
        """Add a single tool. Alias for .tools() with append semantics."""
        self._lists["tools"].append(fn_or_tool)
        return self


    def apply(self, stack: MiddlewareStack) -> Self:
        """Apply a reusable middleware stack (bulk callback registration)."""
        raise NotImplementedError("Implement in hand-written layer")


    def member(self, agent: BaseAgent | AgentBuilder) -> Self:
        """Add a member agent for coordinator pattern."""
        self._lists["sub_agents"].append(agent)
        return self


    def delegate(self, agent) -> 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 guardrail(self, fn: Callable) -> Self:
        """Attach a guardrail function as both before_model and after_model callback."""
        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):
        """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:
        """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

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

    def __getattr__(self, name: str):
        """Forward unknown methods to LlmAgent.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 LlmAgent.model_fields:
            available = sorted(
                set(LlmAgent.model_fields.keys())
                | set(_ALIASES.keys())
                | set(_CALLBACK_ALIASES.keys())
            )
            raise AttributeError(
                f"'{name}' is not a recognized field on LlmAgent. "
                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) -> LlmAgent:
        """LLM-based Agent. Resolve into a native ADK LlmAgent."""
        config = self._prepare_build_config()
        return LlmAgent(**config)
