"""Async sub-namespace for document operations on a Pinecone index.

This module provides the async ``AsyncDocuments`` class that exposes async methods
for upserting and searching flat JSON documents with schema-based indexing.
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

import msgspec
import orjson

from pinecone._internal.constants import DOCUMENTS_HEADERS
from pinecone._internal.introspection import InteractiveIntrospectionMixin
from pinecone._internal.validation import require_in_range, require_non_empty
from pinecone.models.documents.search import Document, DocumentSearchResponse, Usage
from pinecone.models.documents.upsert import DocumentUpsertResponse


if TYPE_CHECKING:
    from httpx import Response

    from pinecone.models.batch import BatchResult


# Decoders for response parsing
_UPSERT_DECODER = msgspec.json.Decoder(DocumentUpsertResponse)


class AsyncDocuments(InteractiveIntrospectionMixin):
    """Async sub-namespace for document operations on a Pinecone index.

    Access via the ``documents`` property on an AsyncIndex instance::

        index = AsyncIndex(host="...", api_key="...")
        response = await index.documents.search(
            namespace="ns",
            top_k=10,
            score_by=[{"type": "text", "field": "title", "text_query": "hello"}],
        )

    Args:
        http: Async HTTP client instance for making requests.
        base_url: Base URL of the index (e.g., ``"https://my-index.svc.pinecone.io"``).
    """

    def __init__(self, *, http: Any, base_url: str) -> None:
        self._http = http
        self._base_url = base_url

    async def _make_data_plane_request(
        self,
        method: str,
        path: str,
        body: dict[str, Any] | None = None,
    ) -> bytes:
        """Make an async data plane request and return raw response bytes.

        Args:
            method: HTTP method (GET, POST, etc.)
            path: Request path
            body: Optional request body (will be JSON-encoded)

        Returns:
            Raw response bytes

        Raises:
            BadRequestError: If the request is invalid (400)
            UnauthorizedError: If authentication fails (401)
            ForbiddenError: If permission denied (403)
            NotFoundError: If resource not found (404)
            RateLimitError: If rate limit exceeded (429)
            ServerError: If server error occurs (5xx)
            APIError: For other HTTP errors
            APIConnectionError: If connection fails
            TimeoutError: If request times out
        """
        url = f"{self._base_url}{path}"
        content = orjson.dumps(body) if body is not None else None

        response: Response = await getattr(self._http, method.lower())(
            url, headers=DOCUMENTS_HEADERS, content=content
        )
        return response.content

    async def upsert(
        self,
        *,
        namespace: str,
        documents: list[dict[str, Any]],
    ) -> DocumentUpsertResponse:
        """Upsert flat JSON documents into a namespace.

        Documents are indexed based on the configured index schema. Vector fields
        can be user-specified (e.g., `my_vector`) or use the reserved `_values`
        key. Text fields are indexed based on schema configuration with
        `full_text_searchable: true`.

        Args:
            namespace: The namespace to upsert documents into. Must not be empty.
            documents: List of flat JSON documents. Each document must have an
                `_id` field. Can contain any fields including metadata, vector
                fields, and text fields. Maximum 1000 documents per request.

        Returns:
            DocumentUpsertResponse with the count of documents upserted.

        Raises:
            ValueError: If namespace is empty, documents list is empty, or
                documents have empty or duplicate _id fields.
            TypeError: If documents have non-string _id fields.
            BadRequestError: If documents are invalid (invalid schema, etc.).
            UnauthorizedError: If authentication fails (401).
            ForbiddenError: If permission denied (403).
            NotFoundError: If namespace not found (404).
            RateLimitError: If rate limit exceeded (429).
            ServerError: If server error occurs (5xx).
            APIError: For other API errors.
            APIConnectionError: If connection fails.
            TimeoutError: If request times out.

        Examples:
            Upsert documents with user-specified vector field:

            >>> response = await index.documents.upsert(
            ...     namespace="my-namespace",
            ...     documents=[
            ...         {
            ...             "_id": "doc1",
            ...             "title": "Hello world",
            ...             "my_vector": [0.1, 0.2, 0.3],
            ...             "genre": "documentary",
            ...         },
            ...         {
            ...             "_id": "doc2",
            ...             "title": "Example",
            ...             "my_vector": [0.4, 0.5, 0.6],
            ...             "genre": "comedy",
            ...         },
            ...     ],
            ... )
            >>> print(response.upserted_count)
            2

            Upsert documents with reserved _values field:

            >>> response = await index.documents.upsert(
            ...     namespace="my-namespace",
            ...     documents=[
            ...         {
            ...             "_id": "doc1",
            ...             "_values": [0.1, 0.2, 0.3],
            ...             "genre": "documentary",
            ...         },
            ...     ],
            ... )
        """
        require_non_empty("namespace", namespace)
        if not documents:
            raise ValueError("documents must not be empty")

        require_in_range("documents", len(documents), 1, 1000)

        # Validate each document has valid _id
        seen_ids: set[str] = set()
        for i, doc in enumerate(documents):
            if "_id" not in doc:
                raise ValueError(
                    f"Document at index {i} missing required '_id' field. "
                    "Each document must have a unique '_id' string field."
                )
            doc_id = doc["_id"]
            if not isinstance(doc_id, str):
                raise TypeError(
                    f"Document at index {i} has non-string '_id': {type(doc_id).__name__}. "
                    "The '_id' field must be a string."
                )
            if not doc_id:
                raise ValueError(
                    f"Document at index {i} has empty '_id'. "
                    "The '_id' field must be a non-empty string."
                )
            if doc_id in seen_ids:
                raise ValueError(
                    f"Document at index {i} has duplicate '_id': {doc_id!r}. "
                    "Each document in a batch must have a unique '_id'."
                )
            seen_ids.add(doc_id)

        # Wrap documents array in a JSON object with "documents" key
        data = await self._make_data_plane_request(
            "POST",
            f"/namespaces/{namespace}/documents/upsert",
            body={"documents": documents},
        )
        return _UPSERT_DECODER.decode(data)

    async def batch_upsert(
        self,
        *,
        namespace: str,
        documents: list[dict[str, Any]],
        batch_size: int = 200,
        max_concurrency: int = 4,
        show_progress: bool = True,
    ) -> BatchResult:
        """Upsert documents in concurrent batches.

        Splits *documents* 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:
            namespace (str): The namespace to upsert documents into.
            documents (list[dict[str, Any]]): List of flat JSON documents.
                Each must have an ``_id`` field. Maximum 1000 documents
                per batch (API limit).
            batch_size (int): Maximum documents per batch request
                (>= 1, default 200).
            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.documents.batch_upsert(
            ...     namespace="my-ns",
            ...     documents=[{"_id": f"doc{i}", "title": f"Title {i}"} for i in range(5_000)],
            ... )
            >>> print(f"Upserted {result.successful_item_count}/{result.total_item_count}")

            Retry failed items:

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

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

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

    async def search(
        self,
        *,
        namespace: str,
        top_k: int,
        score_by: list[dict[str, Any]],
        include_fields: list[str] | str | None = None,
        filter: dict[str, Any] | None = None,  # noqa: A002
    ) -> DocumentSearchResponse:
        """Search documents in a namespace.

        Search documents using text search, vector search, or sparse vector search.
        Results can be filtered by metadata and ranked using the `score_by` parameter.

        For v0, a single query can only be ranked by:
        - Pure text query (multiple terms on the same field)
        - Pure vector query (one field)

        Args:
            namespace: The namespace to search. Must not be empty.
            top_k: The number of top results to return (1-10000).
            score_by: Array of query objects to rank documents. For v0, only one
                query is supported: either a pure text query or a pure vector query
                (dense or sparse).
            include_fields: Fields to include in the response. Use `"*"` to return
                all fields, provide a list of field names for specific fields, or
                None to return IDs only. Default is IDs only.
            filter: A filter expression for filtering documents. Supports standard
                metadata operators ($eq, $ne, $in, $nin, $gt, $gte, $lt, $lte) and  # codespell:ignore nin
                the $text_match operator for required/excluded terms in text fields.

        Returns:
            DocumentSearchResponse with matching documents and usage information.

        Raises:
            ValueError: If namespace is empty, top_k is out of range, or
                score_by is empty.
            BadRequestError: If the request parameters are invalid.
            UnauthorizedError: If authentication fails (401).
            ForbiddenError: If permission denied (403).
            NotFoundError: If namespace not found (404).
            RateLimitError: If rate limit exceeded (429).
            ServerError: If server error occurs (5xx).
            APIError: For other API errors.
            APIConnectionError: If connection fails.
            TimeoutError: If request times out.

        Examples:
            Pure text search:

            >>> response = await index.documents.search(
            ...     namespace="my-namespace",
            ...     top_k=10,
            ...     score_by=[
            ...         {
            ...             "type": "text",
            ...             "field": "title",
            ...             "text_query": 'return "pink panther"',
            ...         }
            ...     ],
            ...     filter={"genre": {"$eq": "comedy"}},
            ...     include_fields=["_id", "genre", "title"],
            ... )
            >>> for doc in response.documents:
            ...     print(f"{doc.id}: {doc.score}")

            Pure vector search:

            >>> response = await index.documents.search(
            ...     namespace="my-namespace",
            ...     top_k=10,
            ...     score_by=[
            ...         {
            ...             "type": "vector",
            ...             "field": "synopsis_vector",
            ...             "values": [0.9, 0.8, 0.7],
            ...         }
            ...     ],
            ...     include_fields="*",  # Return all fields
            ... )

            Sparse vector search:

            >>> response = await index.documents.search(
            ...     namespace="my-namespace",
            ...     top_k=10,
            ...     score_by=[
            ...         {
            ...             "type": "vector",
            ...             "field": "synopsis_vector",
            ...             "sparse_values": {
            ...                 "indices": [23243, 98339],
            ...                 "values": [1.0, 1.0],
            ...             },
            ...         }
            ...     ],
            ... )

            Text match filter with vector search:

            >>> response = await index.documents.search(
            ...     namespace="my-namespace",
            ...     top_k=10,
            ...     score_by=[
            ...         {
            ...             "type": "vector",
            ...             "field": "message_vector",
            ...             "values": [0.9, 0.8, 0.7],
            ...         }
            ...     ],
            ...     filter={
            ...         "service_name": {"$eq": "payment-gateway"},
            ...         "raw_logs": {"$text_match": '+"Connection refused" +"Error 10054"'},
            ...     },
            ... )
        """
        require_non_empty("namespace", namespace)
        require_in_range("top_k", top_k, 1, 10000)

        if not score_by:
            raise ValueError(
                "score_by must not be empty. Provide at least one query object to rank documents. "
                'Example: score_by=[{"type": "text", "field": "title", "text_query": "hello"}]'
            )

        body: dict[str, Any] = {
            "top_k": top_k,
            "score_by": score_by,
        }

        if include_fields is not None:
            body["include_fields"] = include_fields
        if filter is not None:
            body["filter"] = filter

        data = await self._make_data_plane_request(
            "POST",
            f"/namespaces/{namespace}/documents/search",
            body=body,
        )

        # Decode search response and construct Document objects
        raw = orjson.loads(data)
        documents = [Document(doc) for doc in raw.get("documents", [])]
        usage_data = raw.get("usage")
        usage = Usage(read_units=usage_data.get("read_units")) if usage_data else None
        return DocumentSearchResponse(documents=documents, usage=usage)

    def __repr__(self) -> str:
        """Return string representation."""
        return f"AsyncDocuments(base_url={self._base_url!r})"
