import logging
import uuid
import asyncio
from abc import ABC, abstractmethod
from collections.abc import AsyncIterator
from typing import Optional, Any

from aduib_rpc.client import ClientConfig
from aduib_rpc.client.midwares import ClientRequestInterceptor, ClientContext
from aduib_rpc.types import AduibRpcRequest, AduibRpcResponse
from aduib_rpc.server.tasks import (
    TaskSubmitRequest,
    TaskSubmitResponse,
    TaskQueryRequest,
    TaskQueryResponse,
    TaskCancelRequest,
    TaskCancelResponse,
    TaskSubscribeRequest,
    TaskEvent,
)
from aduib_rpc.protocol.v2.health import HealthCheckRequest, HealthCheckResponse
from aduib_rpc.client.transports.base import ClientTransport

try:
    import httpx
    from grpc.aio import Channel
except ImportError:
    httpx = None  # type: ignore
    Channel = None  # type: ignore

logger = logging.getLogger(__name__)


class AduibRpcClient(ABC):
    """Abstract base class for a client."""

    def __init__(
        self,
        middleware: list[ClientRequestInterceptor] | None = None,
    ):
        self._middleware = middleware
        if self._middleware is None:
            self._middleware = []

    @abstractmethod
    async def completion(
        self,
        name: str,
        method: str,
        data: Any = None,
        meta: Optional[dict[str, Any]] = None,
        *,
        context: ClientContext | None = None,
    ) -> AsyncIterator[AduibRpcResponse]:
        """Sends a message to the agent.

        This method handles both streaming and non-streaming (polling) interactions
        based on the client configuration and agent capabilities. It will yield
        events as they are received from the agent.

        Args:
            method: The RPC method to call.
            data: The data to send in the request.
            meta: Optional metadata to include in the request.
            context: The client call context.
        """
        raise NotImplementedError

    @abstractmethod
    async def call(
        self,
        request: AduibRpcRequest,
        *,
        context: ClientContext | None = None,
    ) -> AduibRpcResponse:
        """Unary RPC via AduibRpcService.Call."""
        raise NotImplementedError

    @abstractmethod
    async def call_server_stream(
        self,
        request: AduibRpcRequest,
        *,
        context: ClientContext | None = None,
    ) -> AsyncIterator[AduibRpcResponse]:
        """Server-streaming RPC via AduibRpcService.CallServerStream."""
        raise NotImplementedError

    @abstractmethod
    async def call_client_stream(
        self,
        requests,
        *,
        context: ClientContext | None = None,
    ) -> AduibRpcResponse:
        """Client-streaming RPC."""
        raise NotImplementedError

    @abstractmethod
    async def call_bidirectional(
        self,
        requests,
        *,
        context: ClientContext | None = None,
    ) -> AsyncIterator[AduibRpcResponse]:
        """Bidirectional streaming RPC."""
        raise NotImplementedError

    @abstractmethod
    async def task_submit(
        self,
        request: TaskSubmitRequest,
        *,
        context: ClientContext | None = None,
    ) -> TaskSubmitResponse:
        """Submit a task."""
        raise NotImplementedError

    @abstractmethod
    async def task_query(
        self,
        request: TaskQueryRequest,
        *,
        context: ClientContext | None = None,
    ) -> TaskQueryResponse:
        """Query a task."""
        raise NotImplementedError

    @abstractmethod
    async def task_cancel(
        self,
        request: TaskCancelRequest,
        *,
        context: ClientContext | None = None,
    ) -> TaskCancelResponse:
        """Cancel a task."""
        raise NotImplementedError

    @abstractmethod
    async def task_subscribe(
        self,
        request: TaskSubscribeRequest,
        *,
        context: ClientContext | None = None,
    ) -> AsyncIterator[TaskEvent]:
        """Subscribe to task events."""
        raise NotImplementedError

    @abstractmethod
    async def health_check(
        self,
        request: HealthCheckRequest,
        *,
        context: ClientContext | None = None,
    ) -> HealthCheckResponse:
        """Health check for services."""
        raise NotImplementedError

    @abstractmethod
    async def health_watch(
        self,
        request: HealthCheckRequest,
        *,
        context: ClientContext | None = None,
    ) -> AsyncIterator[HealthCheckResponse]:
        """Watch health status."""
        raise NotImplementedError

    async def add_middleware(
        self,
        middleware: ClientRequestInterceptor,
    ) -> None:
        """Adds a middleware to the client.

        Args:
            middleware: The middleware to add.
        """
        self._middleware.append(middleware)


class BaseAduibRpcClient(AduibRpcClient):
    """Base implementation of the AduibRpc client, containing transport-independent logic."""

    def __init__(
        self,
        config: ClientConfig,
        transport: ClientTransport,
        middleware: list[ClientRequestInterceptor] | None = None,
    ):
        super().__init__(middleware)
        self._config = config
        self._transport = transport

    def _prepare_context(
        self,
        context: ClientContext | None,
        meta: Optional[dict[str, Any]] = None,
    ) -> ClientContext:
        if context is None:
            context = ClientContext()

        context.state["http_kwargs"] = (
            {"headers": meta["headers"]} if meta and "headers" in meta else context.state.get("http_kwargs")
        )
        context.state["security_schema"] = meta.get("security_schema") if meta else context.state.get("security_schema")
        context.state["config"] = self._config
        return context

    async def completion(
        self,
        name: str,
        method: str,
        data: Any = None,
        meta: Optional[dict[str, Any]] = None,
        *,
        context: ClientContext | None = None,
    ) -> AsyncIterator[AduibRpcResponse]:
        """Sends a message to the agent.
        This method handles both streaming and non-streaming (polling) interactions
        based on the client configuration and agent capabilities. It will yield
        events as they are received from the agent.
        Args:
            method: The RPC method to call.
            data: The data to send in the request.
            meta: Optional metadata to include in the request.
            context: The client call context.
        Returns:
            An async iterator yielding `AduibRpcResponse` objects as they are received.
        """
        context = self._prepare_context(context, meta)

        if meta:
            if "stream" in meta:
                logger.warning("The 'stream' meta field is managed by the client configuration and will be overridden.")
            meta["stream"] = str(self._config.streaming).lower()

        request = AduibRpcRequest(name=name, method=method, data=data, meta=meta, id=str(uuid.uuid4()))
        if not self._config.streaming:
            response = await self._transport.completion(request, context=context)
            yield response
            return

        async for response in self._transport.completion_stream(request, context=context):
            yield response

    async def call(
        self,
        request: AduibRpcRequest,
        *,
        context: ClientContext | None = None,
    ) -> AduibRpcResponse:
        meta = request.meta if hasattr(request, "meta") else None
        context = self._prepare_context(context, meta)
        return await self._transport.call(request, context=context)

    async def call_server_stream(
        self,
        request: AduibRpcRequest,
        *,
        context: ClientContext | None = None,
    ) -> AsyncIterator[AduibRpcResponse]:
        meta = request.meta if hasattr(request, "meta") else None
        context = self._prepare_context(context, meta)
        async for response in self._transport.call_server_stream(request, context=context):
            yield response

    async def call_client_stream(
        self,
        requests,
        *,
        context: ClientContext | None = None,
    ) -> AduibRpcResponse:
        context = self._prepare_context(context)
        return await self._transport.call_client_stream(requests, context=context)

    async def call_bidirectional(
        self,
        requests,
        *,
        context: ClientContext | None = None,
    ) -> AsyncIterator[AduibRpcResponse]:
        context = self._prepare_context(context)
        async for response in self._transport.call_bidirectional(requests, context=context):
            yield response

    async def task_submit(
        self,
        request: TaskSubmitRequest,
        *,
        context: ClientContext | None = None,
    ) -> TaskSubmitResponse:
        context = self._prepare_context(context)
        return await self._transport.task_submit(request, context=context)

    async def task_query(
        self,
        request: TaskQueryRequest,
        *,
        context: ClientContext | None = None,
    ) -> TaskQueryResponse:
        context = self._prepare_context(context)
        return await self._transport.task_query(request, context=context)

    async def task_cancel(
        self,
        request: TaskCancelRequest,
        *,
        context: ClientContext | None = None,
    ) -> TaskCancelResponse:
        context = self._prepare_context(context)
        return await self._transport.task_cancel(request, context=context)

    async def task_subscribe(
        self,
        request: TaskSubscribeRequest,
        *,
        context: ClientContext | None = None,
    ) -> AsyncIterator[TaskEvent]:
        context = self._prepare_context(context)
        async for event in self._transport.task_subscribe(request, context=context):
            yield event

    async def health_check(
        self,
        request: HealthCheckRequest,
        *,
        context: ClientContext | None = None,
    ) -> HealthCheckResponse:
        context = self._prepare_context(context)
        return await self._transport.health_check(request, context=context)

    async def health_watch(
        self,
        request: HealthCheckRequest,
        *,
        context: ClientContext | None = None,
    ) -> AsyncIterator[HealthCheckResponse]:
        context = self._prepare_context(context)
        async for status in self._transport.health_watch(request, context=context):
            yield status

    async def close(self) -> None:
        close = getattr(self._transport, "aclose", None)
        if callable(close):
            await close()
            return
        close = getattr(self._transport, "close", None)
        if callable(close):
            result = close()
            if asyncio.iscoroutine(result):
                await result
