"""Async inference namespace implementation for embeddings and reranking.

This module provides the async implementation for the Inference API, enabling
concurrent embedding and reranking operations for agent workflows and async frameworks.
"""

from __future__ import annotations

import asyncio

from typing import TYPE_CHECKING, Any, cast

import msgspec
import orjson

from pinecone.__interface__.inference import AsyncInferenceNamespaceProtocol
from pinecone._internal.adapters.inference_adapter import InferenceAdapter
from pinecone._internal.constants import INFERENCE_API_VERSION
from pinecone._internal.introspection import InteractiveIntrospectionMixin
from pinecone._internal.logging import get_logger
from pinecone._internal.validation import require_in_range, require_non_empty
from pinecone.exceptions import APIConnectionError, APIError
from pinecone.models.inference import (
    EmbedRequest,
    EmbedResponse,
    ListModelsResponse,
    RerankDocument,
    RerankRequest,
    RerankResponse,
)


if TYPE_CHECKING:
    from pinecone.__interface__ import EmbedParameters
    from pinecone.__interface__.inference import (
        EmbeddingResponseProtocol,
        ListModelsResponseProtocol,
        RerankResponseProtocol,
    )
    from pinecone._internal.http import AsyncHTTPClient


logger = get_logger(__name__)


class AsyncInference(InteractiveIntrospectionMixin, AsyncInferenceNamespaceProtocol):
    """Async implementation of inference operations for embeddings and reranking.

    This class enables concurrent API calls for embedding and reranking, making it
    ideal for agent workflows that need to process multiple documents in parallel or
    integrate with async frameworks like FastAPI.

    Args:
        http_client: Async HTTP client for making API requests
        base_url: Base URL for the Inference API

    Example:
        >>> from pinecone import AsyncPinecone
        >>> import asyncio
        >>>
        >>> async def main():
        ...     async with AsyncPinecone(api_key="...") as pc:
        ...         # Generate embeddings
        ...         response = await pc.inference.embed(
        ...             model="multilingual-e5-large",
        ...             texts=["Hello world", "Goodbye world"],
        ...             parameters={"input_type": "passage", "truncate": "END"},
        ...         )
        ...         embeddings = response.embeddings
        ...
        ...         # Rerank documents
        ...         response = await pc.inference.rerank(
        ...             model="bge-reranker-v2-m3",
        ...             query="What is machine learning?",
        ...             documents=["ML is a subset of AI", "Python is a language"],
        ...             top_n=1,
        ...         )
        ...         top_doc = response.results[0]
        >>>
        >>> asyncio.run(main())

    Example with concurrent operations:
        >>> async def embed_batches(pc, text_batches):
        ...     tasks = [
        ...         pc.inference.embed(model="multilingual-e5-large", texts=batch)
        ...         for batch in text_batches
        ...     ]
        ...     return await asyncio.gather(*tasks)
    """

    def __init__(
        self,
        http_client: AsyncHTTPClient,
        base_url: str = "https://api.pinecone.io",
    ):
        """Initialize async Inference namespace.

        Args:
            http_client: Async HTTP client for making API requests
            base_url: Base URL for the Inference API. Defaults to production endpoint.
        """
        self._http = http_client
        self._base_url = base_url
        self._adapter = InferenceAdapter()

    async def _make_inference_request(
        self,
        endpoint: str,
        request: EmbedRequest | RerankRequest | None = None,
        timeout: float | None = None,
        params: dict[str, str] | None = None,
    ) -> dict[str, Any]:
        """Make async HTTP request to inference API with standard headers.

        Args:
            endpoint: API endpoint (e.g., "embed", "rerank", "models")
            request: Optional request model to serialize (for POST requests)
            timeout: Optional request timeout in seconds
            params: Optional query parameters (for GET requests)

        Returns:
            Parsed JSON response as dictionary

        Raises:
            APIError: If the API returns an error response
            ConnectionError: If the request fails to connect
            TimeoutError: If the request times out
            asyncio.CancelledError: If the operation is cancelled
        """
        try:
            headers = {
                "X-Pinecone-Api-Version": INFERENCE_API_VERSION,
            }

            kwargs: dict[str, Any] = {"headers": headers}

            if timeout is not None:
                kwargs["timeout"] = timeout

            if request is not None:
                # POST request with body
                headers["Content-Type"] = "application/json"
                kwargs["content"] = msgspec.json.encode(request)
                response = await self._http.post(f"{self._base_url}/{endpoint}", **kwargs)
            else:
                # GET request with query params
                if params:
                    kwargs["params"] = params
                response = await self._http.get(f"{self._base_url}/{endpoint}", **kwargs)

            return cast("dict[str, Any]", orjson.loads(response.content))
        except asyncio.CancelledError:
            logger.debug("Request cancelled", endpoint=endpoint)
            raise
        except APIConnectionError as e:
            # Network/connection errors - preserve type with added context
            raise APIConnectionError(f"Inference {endpoint} request failed: {e.message}") from e
        except APIError as e:
            # API errors - recreate with same attributes and added context
            raise APIError(
                f"Inference {endpoint} request failed: {e.message}",
                status_code=e.status_code,
                response_body=e.response_body,
            ) from e

    async def list_models(
        self,
        *,
        model_type: str | None = None,
        vector_type: str | None = None,
        timeout: float | None = None,
    ) -> ListModelsResponseProtocol:
        """List available embedding and reranking models asynchronously.

        Args:
            model_type: Optional filter by model type ('embed' or 'rerank')
            vector_type: Optional filter by vector type ('dense' or 'sparse') for embedding models
            timeout: Optional request timeout in seconds

        Returns:
            ListModelsResponse containing:
                - models: List of ModelInfo objects with model details

        Raises:
            APIError: If the API returns an error response
            APIConnectionError: If the request fails to connect
            TimeoutError: If the request times out

        Example:
            >>> async with AsyncPinecone(api_key="...") as pc:
            ...     response = await pc.inference.list_models()
            ...     for model in response.models:
            ...         print(f"{model.model} ({model.type}): {model.short_description}")
            ...         if model.type == "embed" and model.supported_dimensions:
            ...             print(f"  Dimensions: {model.supported_dimensions}")

        Example with filters:
            >>> async with AsyncPinecone(api_key="...") as pc:
            ...     # List only embedding models
            ...     response = await pc.inference.list_models(model_type="embed")
            ...     # List only sparse embedding models
            ...     response = await pc.inference.list_models(model_type="embed", vector_type="sparse")

        Example exploring model parameters:
            >>> async with AsyncPinecone(api_key="...") as pc:
            ...     response = await pc.inference.list_models()
            ...     model = response.models[0]
            ...     print(f"Model: {model.model}")
            ...     print(f"Provider: {model.provider_name}")
            ...     print(f"Max sequence length: {model.max_sequence_length}")
            ...     for param in model.supported_parameters:
            ...         print(f"  {param.parameter} ({param.value_type}): required={param.required}")
        """
        logger.info(
            "Listing available models",
            type_filter=model_type,
            vector_type_filter=vector_type,
        )

        # Build query parameters
        params: dict[str, str] = {}
        if model_type is not None:
            params["type"] = model_type
        if vector_type is not None:
            params["vector_type"] = vector_type

        # Make async HTTP request (GET with query params)
        data: dict[str, Any] = await self._make_inference_request(
            "models", timeout=timeout, params=params if params else None
        )

        # Transform response via adapter

        result: ListModelsResponse = self._adapter.to_list_models_response(data)

        logger.debug(
            "Models listed successfully",
            model_count=len(result.models),
        )

        return result

    async def embed(
        self,
        model: str,
        *,
        texts: list[str],
        parameters: EmbedParameters | dict[str, Any] | None = None,
        timeout: float | None = None,
    ) -> EmbeddingResponseProtocol:
        """Generate embeddings for texts using specified model.

        This async method enables concurrent embedding operations, making it ideal
        for processing large document batches in parallel or integrating with async
        web frameworks.

        Args:
            model: Name of the embedding model to use (e.g., "multilingual-e5-large")
            texts: List of text strings to embed (non-empty)
            parameters: Optional model-specific parameters. See EmbedParameters for
                common options. Additional parameters can be passed as a dict.
            timeout: Optional request timeout in seconds

        Returns:
            EmbeddingResponse containing:
                - embeddings: List of embedding vectors
                - model: Name of the model used
                - usage: Token usage information

        Raises:
            ValueError: If inputs are invalid (empty texts, missing model)
            APIError: If the API returns an error response
            APIConnectionError: If the request fails to connect
            TimeoutError: If the request times out

        Example:
            >>> async with AsyncPinecone(api_key="...") as pc:
            ...     response = await pc.inference.embed(
            ...         model="multilingual-e5-large",
            ...         texts=["Hello world", "Goodbye world"],
            ...         parameters={"input_type": "passage", "truncate": "END"},
            ...     )
            ...     print(f"Generated {len(response.embeddings)} embeddings")
            ...     print(f"Dimension: {len(response.embeddings[0])}")
            ...     print(f"Tokens used: {response.usage.total_tokens}")

        Example with concurrent batches:
            >>> async with AsyncPinecone(api_key="...") as pc:
            ...     batches = [texts[i : i + 100] for i in range(0, len(texts), 100)]
            ...     tasks = [
            ...         pc.inference.embed(model="multilingual-e5-large", texts=batch)
            ...         for batch in batches
            ...     ]
            ...     results = await asyncio.gather(*tasks)

        Example with timeout:
            >>> async with AsyncPinecone(api_key="...") as pc:
            ...     try:
            ...         response = await pc.inference.embed(
            ...             model="multilingual-e5-large",
            ...             texts=["Hello world"],
            ...             timeout=5.0,
            ...         )
            ...     except TimeoutError:
            ...         print("Request timed out after 5 seconds")
        """
        # Validation
        require_non_empty("model", model)
        require_non_empty("texts", texts)

        logger.info(
            "Generating embeddings",
            model=model,
            text_count=len(texts),
            has_parameters=parameters is not None,
        )

        # Build request - convert texts to input format expected by API
        request: EmbedRequest = EmbedRequest(
            model=model,
            inputs=[{"text": text} for text in texts],
            parameters=cast("dict[str, Any] | None", parameters),
        )

        # Make async HTTP request
        data: dict[str, Any] = await self._make_inference_request("embed", request, timeout)

        # Transform response via adapter
        result: EmbedResponse = self._adapter.to_embed_response(data)

        logger.debug(
            "Embeddings generated successfully",
            model=result.model,
            embedding_count=len(result.embeddings),
            total_tokens=result.usage.total_tokens,
        )

        return result

    async def rerank(
        self,
        model: str,
        *,
        query: str,
        documents: list[str],
        top_n: int | None = None,
        return_documents: bool = False,
        parameters: dict[str, Any] | None = None,
        timeout: float | None = None,
    ) -> RerankResponseProtocol:
        """Rerank documents by relevance to a query.

        This async method enables non-blocking reranking operations, useful for
        processing multiple queries concurrently or integrating with async frameworks.

        Args:
                model: Name of the reranking model (e.g., "bge-reranker-v2-m3")
                query: Query text to rank documents against
                documents: List of document texts to rerank (non-empty)
                top_n: Optional number of top results to return
                return_documents: Whether to include document text in results
                parameters: Optional model-specific parameters as a dict. Parameters
                    are entirely model-specific - consult model documentation.
                timeout: Optional request timeout in seconds

        Returns:
            RerankResponse containing:
                - results: List of RerankResult with index, score, document
                - model: Name of the model used
                - usage: Token usage information

        Raises:
            ValueError: If inputs are invalid
            APIError: If the API returns an error response
            APIConnectionError: If the request fails to connect
            TimeoutError: If the request times out

        Example:
                >>> async with AsyncPinecone(api_key="...") as pc:
                ...     response = await pc.inference.rerank(
                ...         model="bge-reranker-v2-m3",
                ...         query="What is machine learning?",
                ...         documents=[
                ...             "Machine learning is a subset of artificial intelligence",
                ...             "Python is a programming language",
                ...             "Deep learning uses neural networks",
                ...         ],
                ...         top_n=2,
                ...     )
                ...     for result in response.results:
                ...         print(f"Document {result.index}: score {result.score:.3f}")

            Example with concurrent queries:
                >>> async with AsyncPinecone(api_key="...") as pc:
                ...     queries = ["query1", "query2", "query3"]
                ...     tasks = [
                ...         pc.inference.rerank(model="bge-reranker-v2-m3", query=q, documents=docs)
                ...         for q in queries
                ...     ]
                ...     results = await asyncio.gather(*tasks)

            Example with timeout:
                >>> async with AsyncPinecone(api_key="...") as pc:
                ...     try:
                ...         response = await pc.inference.rerank(
                ...             model="bge-reranker-v2-m3",
                ...             query="What is AI?",
                ...             documents=["AI is...", "ML is..."],
                ...             timeout=5.0,
                ...         )
                ...     except TimeoutError:
                ...         print("Request timed out after 5 seconds")

        Discovery:
            In interactive environments (REPL, notebooks), use these commands to explore:

            >>> from pinecone import AsyncPinecone
            >>> import asyncio
            >>> async def explore():
            ...     async with AsyncPinecone(api_key="...") as pc:
            ...         # See available methods
            ...         print(dir(pc.inference))  # ['embed', 'rerank']
            ...         # Get detailed help
            ...         help(pc.inference.embed)
            >>> asyncio.run(explore())
        """
        # Validation
        require_non_empty("model", model)
        require_non_empty("query", query)
        require_non_empty("documents", documents)

        if top_n is not None:
            require_in_range("top_n", top_n, 1, len(documents))

        logger.info(
            "Reranking documents",
            model=model,
            document_count=len(documents),
            top_n=top_n,
            return_documents=return_documents,
            has_parameters=parameters is not None,
        )

        # Transform documents to structured format
        doc_objects: list[RerankDocument] = [RerankDocument(text=doc) for doc in documents]

        # Build request
        request: RerankRequest = RerankRequest(
            model=model,
            query=query,
            documents=doc_objects,
            top_n=top_n,
            return_documents=return_documents,
            parameters=parameters,
        )

        # Make async HTTP request
        data: dict[str, Any] = await self._make_inference_request("rerank", request, timeout)

        # Transform response via adapter
        result: RerankResponse = self._adapter.to_rerank_response(data)

        logger.debug(
            "Documents reranked successfully",
            model=result.model,
            result_count=len(result.results),
            total_tokens=result.usage.total_tokens,
        )

        return result

    def __repr__(self) -> str:
        """Return string representation of async inference namespace."""
        return f"AsyncInference(base_url='{self._base_url}')"
