"""Async Index data plane client for vector operations.

This module provides the async client for performing data plane operations
on a specific Pinecone index, including upserting, querying, fetching,
and deleting vectors.
"""

from __future__ import annotations

import os
import warnings

from typing import TYPE_CHECKING, Any

import msgspec
import orjson

from pinecone._internal.constants import DATA_PLANE_HEADERS
from pinecone._internal.vector_normalization import normalize_vectors
from pinecone.models.index.delete import DeleteResponse
from pinecone.models.index.describe_index_stats import DescribeIndexStatsResponse
from pinecone.models.index.fetch import FetchResponse
from pinecone.models.index.fetch_by_metadata import FetchByMetadataResponse
from pinecone.models.index.list_vectors import ListVectorsResponse
from pinecone.models.index.query import QueryResponse
from pinecone.models.index.update import UpdateResponse
from pinecone.models.index.upsert import UpsertResponse


if TYPE_CHECKING:
    from collections.abc import Sequence
    from types import TracebackType

    from pinecone.__interface__.index import Metadata, SparseValues, VectorInput
    from pinecone.async_index.bulk_imports import AsyncBulkImports
    from pinecone.async_index.documents import AsyncDocuments
    from pinecone.async_index.namespaces import AsyncNamespaces
    from pinecone.async_index.records import AsyncRecords
    from pinecone.models.batch import BatchResult


# Module-level decoders for hot-path operations
_UPSERT_DECODER = msgspec.json.Decoder(UpsertResponse)
_QUERY_DECODER = msgspec.json.Decoder(QueryResponse)
_FETCH_DECODER = msgspec.json.Decoder(FetchResponse)
_DELETE_DECODER = msgspec.json.Decoder(DeleteResponse)
_UPDATE_DECODER = msgspec.json.Decoder(UpdateResponse)
_LIST_DECODER = msgspec.json.Decoder(ListVectorsResponse)
_FETCH_BY_METADATA_DECODER = msgspec.json.Decoder(FetchByMetadataResponse)
_DESCRIBE_STATS_DECODER = msgspec.json.Decoder(DescribeIndexStatsResponse)


class AsyncIndex:
    """Async client for data plane operations on a Pinecone index.

    The AsyncIndex client connects to a specific index host and provides
    async methods for vector operations. It can be constructed directly or
    via the ``AsyncPinecone.index()`` factory method.

    Direct construction is recommended when you already know the index host,
    as it avoids a control plane call::

        from pinecone import AsyncIndex

        async with AsyncIndex(host="my-index-abc123.svc.pinecone.io", api_key="...") as index:
            await index.upsert(vectors=[{"id": "v1", "values": [0.1, 0.2]}])
            results = await index.query(vector=[0.1, 0.2], top_k=10)

    Using the factory method on an existing AsyncPinecone client shares
    configuration automatically::

        async with AsyncPinecone(api_key="...", timeout=60) as pc:
            index = pc.index(host="my-index-abc123.svc.pinecone.io")
            results = await index.query(vector=[0.1, 0.2], top_k=10)

    Args:
        host: Index host URL for data plane operations.
        api_key: Pinecone API key. If not provided, reads from the
            PINECONE_API_KEY environment variable.
        timeout: Request timeout in seconds. Can be a single number or a dict
            with keys 'connect', 'read', 'write', 'pool'. Defaults to 30.
        additional_headers: Extra headers to include in every request.

    Raises:
        ValueError: If api_key is not provided and PINECONE_API_KEY
            environment variable is not set.

    Example:
        >>> from pinecone import AsyncIndex
        >>> import asyncio
        >>>
        >>> async def main():
        ...     async with AsyncIndex(
        ...         host="my-index-abc123.svc.pinecone.io", api_key="..."
        ...     ) as index:
        ...         await index.upsert(vectors=[{"id": "v1", "values": [0.1, 0.2, 0.3]}])
        ...         results = await index.query(vector=[0.1, 0.2, 0.3], top_k=10)
        ...         for match in results.matches:
        ...             print(f"{match.id}: {match.score}")
        >>>
        >>> asyncio.run(main())
    """

    def __init__(
        self,
        *,
        host: str,
        api_key: str | None = None,
        timeout: int | float | dict[str, float] = 30,
        additional_headers: dict[str, str] | None = None,
    ):
        resolved_api_key = api_key or os.environ.get("PINECONE_API_KEY")
        if not resolved_api_key:
            raise ValueError(
                "api_key is required. Provide it as a parameter or set "
                "the PINECONE_API_KEY environment variable."
            )

        self._host = host
        self._base_url = f"https://{host}"
        self._api_key = resolved_api_key
        self._timeout = timeout
        self._additional_headers = additional_headers

        # Lazy sub-namespace instances
        self._bulk_imports: AsyncBulkImports | None = None
        self._namespaces: AsyncNamespaces | None = None
        self._records: AsyncRecords | None = None
        self._documents: AsyncDocuments | None = None

        from pinecone._internal.http import AsyncHTTPClient, HTTPConfig

        config = HTTPConfig(timeout=timeout)
        self._http = AsyncHTTPClient(api_key=resolved_api_key, config=config)

    async def _make_data_plane_request(
        self,
        method: str,
        path: str,
        *,
        body: dict[str, Any] | None = None,
        params: dict[str, Any] | None = None,
    ) -> bytes:
        """Make an async HTTP request to the data plane API."""
        url = f"{self._base_url}{path}"
        content = orjson.dumps(body) if body is not None else None

        if method == "GET":
            response = await self._http.get(url, headers=DATA_PLANE_HEADERS, params=params)
        else:
            response = await self._http.post(
                url, headers=DATA_PLANE_HEADERS, content=content, params=params
            )

        return response.content

    async def upsert(
        self,
        vectors: Sequence[VectorInput],
        *,
        namespace: str = "",
    ) -> UpsertResponse:
        """Insert or update vectors.

        Args:
            vectors: Sequence of vectors to upsert. Each vector can be:
                - Dict with 'id', 'values', and optionally 'sparseValues' and 'metadata'
                - 2-tuple of (id, values) for simple vectors
                - 3-tuple of (id, values, metadata) for vectors with metadata
            namespace: Target namespace.

        Returns:
            UpsertResponse with the count of upserted vectors.

        Note:
            Dict format is recommended for best performance and sparse vector support.
            Tuple formats are provided for compatibility with prior versions of the SDK.

        Example:
            Dict format (recommended):

            >>> response = await index.upsert(
            ...     vectors=[
            ...         {"id": "v1", "values": [0.1, 0.2], "metadata": {"key": "val"}},
            ...         {"id": "v2", "values": [0.3, 0.4]},
            ...     ],
            ... )

            Tuple formats (compatibility):

            >>> response = await index.upsert(
            ...     vectors=[
            ...         ("v1", [0.1, 0.2], {"key": "val"}),  # 3-tuple
            ...         ("v2", [0.3, 0.4]),  # 2-tuple
            ...     ],
            ... )
            >>> print(f"Upserted {response.upserted_count} vectors")
        """
        # Normalize all vectors to dict format
        normalized_vectors = normalize_vectors(vectors)

        body: dict[str, Any] = {"vectors": normalized_vectors}
        if namespace:
            body["namespace"] = namespace

        data = await self._make_data_plane_request("POST", "/vectors/upsert", body=body)
        return _UPSERT_DECODER.decode(data)

    async def batch_upsert(
        self,
        vectors: Sequence[VectorInput],
        *,
        namespace: str = "",
        batch_size: int = 100,
        max_concurrency: int = 4,
        show_progress: bool = True,
    ) -> BatchResult:
        """Upsert vectors in concurrent batches.

        Splits *vectors* into chunks of *batch_size* and upserts them
        concurrently.  Returns a summary with success/failure counts and
        the failed items for easy retry.

        Individual batch failures are captured in the result rather than
        raised, so a single bad batch does not abort the entire operation.

        Args:
            vectors: Sequence of vectors to upsert. Each vector can be:
                - Dict with 'id', 'values', and optionally 'sparseValues' and 'metadata'
                - 2-tuple of (id, values) for simple vectors
                - 3-tuple of (id, values, metadata) for vectors with metadata
            namespace (str): Target namespace (default ``""``).
            batch_size (int): Maximum vectors per batch request (>= 1).
            max_concurrency (int): Maximum concurrent batch requests (1-64).
            show_progress (bool): Show a tqdm progress bar (if installed).

        Returns:
            BatchResult with counts and any errors for retry.

        Raises:
            ValueError: If *batch_size* or *max_concurrency* is out of range.

        Examples:
            >>> result = await index.batch_upsert(
            ...     vectors=[{"id": f"v{i}", "values": [0.1] * 128} for i in range(10_000)],
            ...     namespace="my-ns",
            ...     batch_size=100,
            ... )
            >>> print(f"Upserted {result.successful_item_count}/{result.total_item_count}")

            Retry failed items:

            >>> if result.has_errors:
            ...     retry = await index.batch_upsert(vectors=result.failed_items, namespace="my-ns")
        """
        from pinecone._internal.batch import async_batch_execute

        # Normalize all vectors to dict format for batch processing
        normalized_vectors = normalize_vectors(vectors)

        async def _upsert_batch(batch: list[dict[str, Any]]) -> UpsertResponse:
            return await self.upsert(vectors=batch, namespace=namespace)

        return await async_batch_execute(
            items=normalized_vectors,
            operation=_upsert_batch,
            batch_size=batch_size,
            max_concurrency=max_concurrency,
            show_progress=show_progress,
            desc="Upserting vectors",
        )

    async def query(
        self,
        *,
        vector: list[float] | None = None,
        id: str | None = None,  # noqa: A002
        top_k: int = 10,
        namespace: str = "",
        filter: dict[str, Any] | None = None,  # noqa: A002
        include_values: bool = False,
        include_metadata: bool = True,
        sparse_vector: SparseValues | None = None,
    ) -> QueryResponse:
        """Query vectors by similarity.

        Args:
            vector: Dense query vector.
            id: Query by existing vector ID.
            top_k: Number of results to return.
            namespace: Target namespace.
            filter: Metadata filter expression.
            include_values: Include vector values in results.
            include_metadata: Include metadata in results.
            sparse_vector: Optional sparse query vector.

        Returns:
            QueryResponse with matching vectors and scores.

        Example:
            >>> results = await index.query(vector=[0.1, 0.2, 0.3], top_k=10)
        """
        body: dict[str, Any] = {
            "topK": top_k,
            "includeValues": include_values,
            "includeMetadata": include_metadata,
        }
        if vector is not None:
            body["vector"] = vector
        if id is not None:
            body["id"] = id
        if namespace:
            body["namespace"] = namespace
        if filter is not None:
            body["filter"] = filter
        if sparse_vector is not None:
            body["sparseVector"] = sparse_vector

        data = await self._make_data_plane_request("POST", "/query", body=body)
        return _QUERY_DECODER.decode(data)

    async def fetch(
        self,
        ids: list[str],
        *,
        namespace: str = "",
    ) -> FetchResponse:
        """Fetch vectors by ID.

        Args:
            ids: List of vector IDs to fetch.
            namespace: Namespace to fetch from.

        Returns:
            FetchResponse with the fetched vectors.

        Example:
            >>> response = await index.fetch(ids=["vec1", "vec2"])
        """
        params: dict[str, Any] = {"ids": ids}
        if namespace:
            params["namespace"] = namespace

        data = await self._make_data_plane_request("GET", "/vectors/fetch", params=params)
        return _FETCH_DECODER.decode(data)

    async def delete(
        self,
        *,
        ids: list[str] | None = None,
        delete_all: bool = False,
        namespace: str = "",
        filter: dict[str, Any] | None = None,  # noqa: A002
    ) -> DeleteResponse:
        """Delete vectors.

        Args:
            ids: Vector IDs to delete.
            delete_all: Delete all vectors in the namespace.
            namespace: Target namespace.
            filter: Delete vectors matching this metadata filter.

        Returns:
            DeleteResponse (empty on success).

        Example:
            >>> await index.delete(ids=["vec1", "vec2"])
        """
        body: dict[str, Any] = {}
        if ids is not None:
            body["ids"] = ids
        if delete_all:
            body["deleteAll"] = True
        if namespace:
            body["namespace"] = namespace
        if filter is not None:
            body["filter"] = filter

        data = await self._make_data_plane_request("POST", "/vectors/delete", body=body)
        return _DELETE_DECODER.decode(data)

    async def update(
        self,
        id: str,  # noqa: A002
        *,
        values: list[float] | None = None,
        sparse_values: SparseValues | None = None,
        set_metadata: Metadata | None = None,
        namespace: str = "",
    ) -> UpdateResponse:
        """Update a vector by ID.

        Args:
            id: Vector ID to update.
            values: New dense vector values.
            sparse_values: New sparse vector values.
            set_metadata: Metadata fields to set or overwrite.
            namespace: Target namespace.

        Returns:
            UpdateResponse with the number of matched records.

        Example:
            >>> response = await index.update(id="vec1", set_metadata={"k": "v"})
        """
        body: dict[str, Any] = {"id": id}
        if values is not None:
            body["values"] = values
        if sparse_values is not None:
            body["sparseValues"] = sparse_values
        if set_metadata is not None:
            body["setMetadata"] = set_metadata
        if namespace:
            body["namespace"] = namespace

        data = await self._make_data_plane_request("POST", "/vectors/update", body=body)
        return _UPDATE_DECODER.decode(data)

    async def list(
        self,
        *,
        prefix: str | None = None,
        limit: int | None = None,
        pagination_token: str | None = None,
        namespace: str = "",
    ) -> ListVectorsResponse:
        """List vector IDs in a namespace.

        Args:
            prefix: Filter to IDs starting with this prefix.
            limit: Maximum number of IDs to return.
            pagination_token: Token to resume from a previous call.
            namespace: Namespace to list from.

        Returns:
            ListVectorsResponse with vector IDs and pagination info.

        Example:
            >>> response = await index.list(prefix="doc#", limit=100)
        """
        params: dict[str, Any] = {}
        if prefix is not None:
            params["prefix"] = prefix
        if limit is not None:
            params["limit"] = limit
        if pagination_token is not None:
            params["paginationToken"] = pagination_token
        if namespace:
            params["namespace"] = namespace

        data = await self._make_data_plane_request("GET", "/vectors/list", params=params)
        return _LIST_DECODER.decode(data)

    async def fetch_by_metadata(
        self,
        *,
        filter: dict[str, Any] | None = None,  # noqa: A002
        namespace: str = "",
        limit: int | None = None,
        pagination_token: str | None = None,
    ) -> FetchByMetadataResponse:
        """Fetch vectors matching a metadata filter.

        Args:
            filter: Metadata filter expression.
            namespace: Namespace to fetch from.
            limit: Maximum number of vectors to return.
            pagination_token: Token to resume from a previous call.

        Returns:
            FetchByMetadataResponse with matching vectors.

        Example:
            >>> response = await index.fetch_by_metadata(filter={"genre": "comedy"})
        """
        body: dict[str, Any] = {}
        if filter is not None:
            body["filter"] = filter
        if namespace:
            body["namespace"] = namespace
        if limit is not None:
            body["limit"] = limit
        if pagination_token is not None:
            body["paginationToken"] = pagination_token

        data = await self._make_data_plane_request("POST", "/vectors/fetch_by_metadata", body=body)
        return _FETCH_BY_METADATA_DECODER.decode(data)

    async def describe_index_stats(
        self,
        *,
        filter: dict[str, Any] | None = None,  # noqa: A002
    ) -> DescribeIndexStatsResponse:
        """Get index statistics.

        Args:
            filter: Optional metadata filter for filtered stats.

        Returns:
            DescribeIndexStatsResponse with index statistics.

        Example:
            >>> stats = await index.describe_index_stats()
            >>> print(f"Total vectors: {stats.total_vector_count}")
        """
        body: dict[str, Any] = {}
        if filter is not None:
            body["filter"] = filter

        data = await self._make_data_plane_request("POST", "/describe_index_stats", body=body)
        return _DESCRIBE_STATS_DECODER.decode(data)

    @property
    def bulk_imports(self) -> AsyncBulkImports:
        """Sub-namespace for bulk import operations.

        Example:
            >>> imports = await index.bulk_imports.list()
        """
        if self._bulk_imports is None:
            from pinecone.async_index.bulk_imports import AsyncBulkImports

            self._bulk_imports = AsyncBulkImports(http=self._http, base_url=self._base_url)
        return self._bulk_imports

    @property
    def namespaces(self) -> AsyncNamespaces:
        """Sub-namespace for namespace management operations.

        Example:
            >>> ns_list = await index.namespaces.list()
        """
        if self._namespaces is None:
            from pinecone.async_index.namespaces import AsyncNamespaces

            self._namespaces = AsyncNamespaces(http=self._http, base_url=self._base_url)
        return self._namespaces

    @property
    def records(self) -> AsyncRecords:
        """Sub-namespace for records API operations.

        Example:
            >>> results = await index.records.search(
            ...     namespace="ns", query={"top_k": 10, "inputs": {"text": "hello"}}
            ... )
        """
        if self._records is None:
            from pinecone.async_index.records import AsyncRecords

            self._records = AsyncRecords(http=self._http, base_url=self._base_url)
        return self._records

    @property
    def documents(self) -> AsyncDocuments:
        """Access async document operations namespace.

        The AsyncDocuments namespace provides async methods for upserting and
        searching flat JSON documents with schema-based indexing.

        Returns:
            AsyncDocuments namespace instance.

        Example:
            >>> response = await index.documents.search(
            ...     namespace="my-ns",
            ...     top_k=10,
            ...     score_by=[{"type": "text", "field": "title", "text_query": "hello"}],
            ... )
        """
        if self._documents is None:
            from pinecone.async_index.documents import AsyncDocuments

            self._documents = AsyncDocuments(http=self._http, base_url=self._base_url)
        return self._documents

    async def close(self) -> None:
        """Close the client and release connections.

        After calling close(), the client should not be used.

        Example:
            >>> index = AsyncIndex(host="...", api_key="...")
            >>> try:
            ...     results = await index.query(vector=[0.1, 0.2], top_k=5)
            ... finally:
            ...     await index.close()
        """
        await self._http.close()

    async def __aenter__(self) -> AsyncIndex:
        """Enable async context manager usage.

        Example:
            >>> async with AsyncIndex(host="...", api_key="...") as index:
            ...     results = await index.query(vector=[0.1, 0.2], top_k=5)
        """
        return self

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> None:
        """Clean up when exiting async context manager."""
        await self.close()

    def __repr__(self) -> str:
        """Return string representation."""
        return f"AsyncIndex(host='{self._host}')"

    def __del__(self) -> None:
        """Warn if client wasn't properly closed."""
        try:
            if self._http._client is not None and not self._http._client.is_closed:
                warnings.warn(
                    f"Unclosed {self.__class__.__name__}. "
                    "Use 'async with' or call 'await index.close()'",
                    ResourceWarning,
                    stacklevel=2,
                )
        except Exception:  # noqa: BLE001, S110
            pass
