from typing import Any, Callable, Optional, List, Dict

class Agent:
    def __init__(
        self,
        *,
        agent_id: Optional[str] = ...,
        auto_save: bool = ...,
        parent_agent_runtime_ctx: Any = ...,
        parent_tool_manager: Any = ...,
        global_storage: Any = ...,
        parent_plugin_manager: Any = ...,
        parent_settings: Any = ...,
        is_debug: Optional[bool] = ...,
    ): ...
    def refresh_plugins(self) -> None: ...
    def toggle_auto_save(self, is_enabled: bool) -> "Agent": ...
    def toggle_component(self, component_name: str, is_enabled: bool) -> "Agent": ...
    def save(self) -> "Agent": ...
    def set_settings(self, settings_key: str, settings_value: Any) -> "Agent": ...
    def set_global_variable(
        self, variable_name: str, variable_value: Any
    ) -> "Agent": ...
    def set_agent_prompt(self, key: str, value: Any) -> "Agent": ...
    def get_agent_prompt(self, key: str) -> Any: ...
    def remove_agent_prompt(self, key: str) -> "Agent": ...
    def register_agent_component(
        self, agent_component_name: str, agent_component_class: Callable
    ) -> "Agent": ...
    def set_agent_component_settings(
        self, agent_component_name: str, key: str, value: Any
    ) -> "Agent": ...
    def attach_workflow(self, name: str, workflow: Any) -> "Agent": ...
    async def start_async(
        self, request_type: Optional[str] = ..., *, return_generator: bool = ...
    ) -> Any: ...
    def start(
        self, request_type: Optional[str] = ..., *, return_generator: bool = ...
    ) -> Any: ...

    # ==== alias_manager 动态注入的常用方法 ====
    def general(self, *args, **kwargs) -> "Agent": ...
    def role(self, *args, **kwargs) -> "Agent": ...
    def user_info(self, *args, **kwargs) -> "Agent": ...
    def abstract(self, *args, **kwargs) -> "Agent": ...
    def chat_history(self, list) -> "Agent": ...
    def input(self, *args, **kwargs) -> "Agent": ...
    def info(self, *args, **kwargs) -> "Agent": ...
    def instruct(self, *args, **kwargs) -> "Agent": ...
    def output(self, *args, **kwargs) -> "Agent": ...
    def file(self, *args, **kwargs) -> "Agent": ...
    def files(self, *args, **kwargs) -> "Agent": ...
    def set_request_prompt(self, key: str, value: Any) -> "Agent": ...
    def get_request_prompt(self, key: str) -> "Agent": ...
    def remove_request_prompt(self, key: str) -> "Agent": ...

    # ==== plugins/agent_component 下所有 alias ====
    # Decorator
    def auto_func(self, func: Callable) -> Callable: ...
    def on_event(self, event: str = ..., *, is_await: bool = ...) -> Callable: ...
    def tool(self, **tool_info_kwrags) -> Callable: ...

    # EventListener
    def add_event_listener(
        self,
        event: str,
        listener: Callable,
        *,
        is_await: bool = ...,
        is_agent_event: bool = ...,
    ) -> "Agent": ...
    def on_delta(
        self, listener: Callable, *, is_await: bool = ..., is_agent_event: bool = ...
    ) -> "Agent": ...
    def on_done(
        self, listener: Callable, *, is_await: bool = ..., is_agent_event: bool = ...
    ) -> "Agent": ...
    def on_finally(
        self, listener: Callable, *, is_await: bool = ..., is_agent_event: bool = ...
    ) -> "Agent": ...
    def on_instant(
        self, listener: Callable, *, is_await: bool = ..., is_agent_event: bool = ...
    ) -> "Agent": ...
    def call_event_listeners(self, event: str, data: Any) -> Any: ...

    # Instant
    def use_instant(self) -> "Agent": ...
    def use_realtime(self) -> "Agent": ...

    # ReplyReformer
    def reform_reply(self, reform_handler: Callable) -> "Agent": ...

    # ResponseGenerator
    def put_data_to_generator(self, event: Any, data: Any) -> None: ...
    def get_generator(self) -> Any: ...
    def get_delta_generator(self) -> Any: ...
    def get_instant_generator(self) -> Any: ...
    def get_realtime_generator(self) -> Any: ...
    def get_instant_keys_generator(self, keys: Any) -> Any: ...
    def get_complete_generator(self) -> Any: ...

    # Role
    def set_role_id(self, role_id: str) -> "Agent": ...
    def set_role(self, key: Any, value: Any = ...) -> "Agent": ...
    def update_role(self, key: Any, value: Any = ...) -> "Agent": ...
    def append_role(self, key: Any, value: Any = ...) -> "Agent": ...
    def extend_role(self, key: Any, value: Any = ...) -> "Agent": ...
    def save_role(self, role_id: str = ...) -> "Agent": ...
    def load_role(self, role_id: str) -> "Agent": ...

    # Search
    def use_search(self) -> "Agent": ...

    # Segment
    def segment(
        self,
        name: str,
        output_prompt: Any,
        listener: Callable = ...,
        *,
        is_streaming: bool = ...,
        is_await: bool = ...,
    ) -> "Agent": ...
    def add_segment_listener(
        self,
        name: str,
        listener: Callable,
        *,
        is_streaming: bool = ...,
        is_await: bool = ...,
    ) -> "Agent": ...
    def on_segment_delta(
        self, name: str, listener: Callable, *, is_await: bool = ...
    ) -> "Agent": ...
    def on_segment_done(
        self, name: str, listener: Callable, *, is_await: bool = ...
    ) -> "Agent": ...

    # Session
    def active_session(self, session_id: str = ...) -> str: ...
    def activate_session(self, session_id: str = ...) -> str: ...
    def save_session(self) -> "Agent": ...
    def stop_session(self) -> "Agent": ...
    def toggle_session_auto_save(self, is_enabled: bool) -> "Agent": ...
    def toggle_strict_orders(self, is_strict_orders: bool) -> "Agent": ...
    def toggle_manual_chat_history(self, is_manual_chat_history: bool) -> "Agent": ...
    def set_chat_history_max_length(self, max_length: int) -> "Agent": ...
    def add_chat_history(self, role: str, content: str) -> "Agent": ...
    def get_chat_history(self, *, is_shorten: bool = ...) -> Any: ...
    def rewrite_chat_history(self, new_chat_history: list) -> "Agent": ...
    def set_abstract(self, *args, **kwargs) -> "Agent": ...

    # Status
    def set_status(self, key: str, value: str) -> "Agent": ...
    def save_status(self) -> "Agent": ...
    def load_status(self) -> "Agent": ...
    def use_global_status(self, namespace_name: str = ...) -> "Agent": ...
    def append_status_mapping(
        self, status_key: str, status_value: str, alias_name: str, *args, **kwargs
    ) -> "Agent": ...

    # Tool
    def register_tool(
        self,
        tool_name: str,
        desc: Any,
        args: dict,
        func: Callable,
        *,
        categories: Any = ...,
    ) -> "Agent": ...
    def call_tool(self, *args, **kwargs) -> Any: ...
    def set_tool_proxy(self, *args, **kwargs) -> "Agent": ...
    def add_public_tools(self, tool_names: Any) -> "Agent": ...
    def use_public_tools(self, tool_names: Any) -> "Agent": ...
    def add_public_categories(self, tool_categories: Any) -> "Agent": ...
    def use_public_categories(self, tool_categories: Any) -> "Agent": ...
    def add_all_public_tools(self) -> "Agent": ...
    def use_all_public_tools(self) -> "Agent": ...
    def use_mcp_server(
        self,
        command: str = ...,
        args: List[str] = ...,
        env: str = ...,
        *,
        config: dict = ...,
        categories: Any = ...,
    ) -> "Agent": ...
    def set_tool_using_workflow(self, workflow: Any) -> "Agent": ...
    def remove_tools(self, tool_names: Any) -> "Agent": ...
    def must_call(self, tool_name: str) -> "Agent": ...

    # UserInfo
    def set_user_name(self, name: str) -> "Agent": ...
    def set_user_info(self, key: Any, value: Any = ...) -> "Agent": ...
    def update_user_info(self, key: Any, value: Any = ...) -> "Agent": ...
    def append_user_info(self, key: Any, value: Any = ...) -> "Agent": ...
    def extend_user_info(self, key: Any, value: Any = ...) -> "Agent": ...
    def save_user_info(self, role_name: str = ...) -> "Agent": ...
    def load_user_info(self, role_name: str) -> "Agent": ...

    # YAMLLoader
    def load_yaml_prompt(
        self,
        *,
        path: str = ...,
        yaml: str = ...,
        use_agently_style: bool = ...,
        variables: dict = ...,
    ) -> "Agent": ...
