import aiohttp
from tritonclient.http.aio import InferenceServerClient as AsyncInferenceServerClient

from llm_common.clients.aiohttp_client import ClientSessionWithMonitoring


class TritonHTTPClientSessionWithMonitoring(ClientSessionWithMonitoring):
    """AIOHTTP session for Triton with monitoring enabled."""

    name_for_monitoring = "triton"

    def clear_resource_path(self, resource: str) -> str:
        """Clean resource path for metrics grouping."""
        return resource


class AsyncInferenceServerClientWithMonitoring(AsyncInferenceServerClient):
    """Async Triton client that uses ClientSessionWithMonitoring for HTTP metrics.

    This subclass overrides __init__ to inject our monitored aiohttp session
    instead of the default one.
    """

    def __init__(
        self,
        url,
        verbose=False,
        conn_limit=100,
        conn_timeout=60.0,
        ssl=False,
        ssl_context=None,
    ):
        super().__init__(
            url, verbose=verbose, conn_limit=conn_limit, conn_timeout=conn_timeout, ssl=ssl, ssl_context=ssl_context
        )
        self._stub = TritonHTTPClientSessionWithMonitoring(
            connector=self._conn,
            timeout=aiohttp.ClientTimeout(total=conn_timeout),
            auto_decompress=self._stub._auto_decompress,
        )
