# -*- coding: utf-8 -*-
# filename: base_client.pyi
# @Time    : 2025/8/19 10:28
# @Author  : JQQ
# @Email   : jqq1716@gmail.com
# @Software: PyCharm
import asyncio
from abc import ABC, abstractmethod
from collections.abc import Awaitable
from collections.abc import Callable as Callable
from contextlib import AsyncExitStack
from enum import StrEnum
from typing import Any, Generic, TypeVar

from mcp import ClientSession
from mcp import Tool as Tool
from mcp.client.session import MessageHandlerFnT
from mcp.types import CallToolResult, InitializeResult, ReadResourceResult, Resource
from pydantic import BaseModel as BaseModel
from transitions.core import EventData
from transitions.extensions.asyncio import AsyncMachine, AsyncTransitionConfigDict

from a2c_smcp.utils.async_property import async_property

# 泛型参数，用于约束 MCP Server 参数类型
# Generic parameter for constraining MCP Server parameter types
ParamsT = TypeVar("ParamsT", bound=BaseModel)

class STATES(StrEnum):
    initialized = "initialized"
    connected = "connected"
    disconnected = "disconnected"
    error = "error"

TRANSITIONS: list[AsyncTransitionConfigDict]

class A2CAsyncMachine(AsyncMachine): ...

class BaseMCPClient(ABC, Generic[ParamsT]):
    params: ParamsT
    _aexit_stack: AsyncExitStack
    machine: A2CAsyncMachine
    _async_session: ClientSession | None
    _state_change_callback: Callable[[str, str], None | Awaitable[None]] | None
    _message_handler: MessageHandlerFnT | None
    _session_keep_alive_task: asyncio.Task | None
    _create_session_success_event: asyncio.Event
    _create_session_failure_event: asyncio.Event
    _async_session_closed_event: asyncio.Event
    _initialize_result: InitializeResult | None
    state: STATES

    def __init__(
        self,
        params: ParamsT,
        state_change_callback: Callable[[str, str], None | Awaitable[None]] | None = None,
        message_handler: MessageHandlerFnT | None = None,
    ) -> None: ...
    async def _trigger_state_change(self, event: EventData) -> None: ...
    @async_property
    async def async_session(self) -> ClientSession: ...
    @abstractmethod
    async def _create_async_session(self) -> ClientSession: ...
    @property
    def initialize_result(self) -> InitializeResult | None: ...
    async def aconnect(self, *args: Any, **kwargs: Any) -> None: ...
    async def aprepare_connect(self, event: EventData) -> None: ...
    async def acan_connect(self, event: EventData) -> bool: ...
    async def abefore_connect(self, event: EventData) -> None: ...
    async def on_enter_connected(self, event: EventData) -> None: ...
    async def aafter_connect(self, event: EventData) -> None: ...
    async def adisconnect(self, *args: Any, **kwargs: Any) -> None: ...
    async def aprepare_disconnect(self, event: EventData) -> None: ...
    async def acan_disconnect(self, event: EventData) -> bool: ...
    async def abefore_disconnect(self, event: EventData) -> None: ...
    async def on_enter_disconnected(self, event: EventData) -> None: ...
    async def aafter_disconnect(self, event: EventData) -> None: ...
    async def aerror(self, *args: Any, **kwargs: Any) -> None: ...
    async def aprepare_error(self, event: EventData) -> None: ...
    async def acan_error(self, event: EventData) -> bool: ...
    async def abefore_error(self, event: EventData) -> None: ...
    async def on_enter_error(self, event: EventData) -> None: ...
    async def aafter_error(self, event: EventData) -> None: ...
    async def ainitialize(self, *args: Any, **kwargs: Any) -> None: ...
    async def aprepare_initialize(self, event: EventData) -> None: ...
    async def acan_initialize(self, event: EventData) -> bool: ...
    async def abefore_initialize(self, event: EventData) -> None: ...
    async def on_enter_initialized(self, event: EventData) -> None: ...
    async def aafter_initialize(self, event: EventData) -> None: ...
    async def list_tools(self) -> list[Tool]: ...
    async def list_windows(self) -> list[Resource]: ...
    async def get_window_detail(self, resource: Resource | str) -> ReadResourceResult: ...
    async def call_tool(self, tool_name: str, params: dict) -> CallToolResult: ...
    async def _keep_alive_task(self) -> None: ...
    async def _close_task(self) -> None: ...
