"""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 msgspec import Struct

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
    from pinecone.models.documents.score_by import ScoreByQuery


# 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. Each document
        must have an `_id` field. Documents can contain:
        - Dense vectors via `_values` field
        - Sparse vectors via `_sparse_values` field with `indices` and `values` arrays
        - Text and metadata fields for full-text search
        - String array values for metadata (e.g., tags)

        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. Maximum 100 documents per request.

        Returns:
            DocumentUpsertResponse with the count of documents upserted.

        Raises:
            ValueError: If namespace is empty, documents list is empty,
                documents count exceeds 100, or documents have empty or
                duplicate _id fields.
            TypeError: If documents have non-string _id fields.
            BadRequestError: If documents are invalid (invalid schema, dimension
                mismatch, 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:
            Dense vector upsert:

            >>> response = await index.documents.upsert(
            ...     namespace="movies",
            ...     documents=[
            ...         {
            ...             "_id": "movie1",
            ...             "text": "A sci-fi epic about interstellar travel",
            ...             "genre": "sci-fi",
            ...             "tags": ["space", "adventure"],
            ...             "_values": [0.1, 0.2, 0.3],
            ...         },
            ...         {
            ...             "_id": "movie2",
            ...             "text": "A romantic comedy set in New York",
            ...             "genre": "comedy",
            ...             "tags": ["romance", "funny"],
            ...             "_values": [0.4, 0.5, 0.6],
            ...         },
            ...     ],
            ... )
            >>> print(response.upserted_count)
            2

            Sparse vector upsert:

            >>> response = await index.documents.upsert(
            ...     namespace="movies",
            ...     documents=[
            ...         {
            ...             "_id": "movie1",
            ...             "text": "A sci-fi epic about interstellar travel",
            ...             "genre": "sci-fi",
            ...             "_sparse_values": {
            ...                 "indices": [0, 5, 10, 100],
            ...                 "values": [0.5, 0.3, 0.8, 0.2],
            ...             },
            ...         },
            ...     ],
            ... )
            >>> print(response.upserted_count)
            1

            Text-only document (for full-text search indexes):

            >>> response = await index.documents.upsert(
            ...     namespace="movies",
            ...     documents=[
            ...         {
            ...             "_id": "movie1",
            ...             "text": "A sci-fi epic about interstellar travel",
            ...             "genre": "sci-fi",
            ...         },
            ...     ],
            ... )
            >>> print(response.upserted_count)
            1
        """
        require_non_empty("namespace", namespace)
        if not documents:
            raise ValueError("documents must not be empty")

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

        # 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 = 100,
        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 100 documents
                per batch (API limit).
            batch_size (int): Maximum documents per batch request
                (>= 1, <= 100, default 100).
            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] | ScoreByQuery],
        include_fields: list[str] | str | None = None,
        filter: dict[str, Any] | None = None,  # noqa: A002
    ) -> DocumentSearchResponse:
        """Search documents in a namespace.

        Finds documents matching the provided query criteria. Supports dense vectors,
        sparse vectors, full-text search, and query string search. Results are ranked
        by the scoring method specified in score_by.

        Args:
            namespace: Namespace to search within. Must not be empty.
            top_k: Maximum number of documents to return (1-10000).
            score_by: List of query specifications. Each query specifies a scoring
                method (dense_vector, sparse_vector, text, or query_string) and
                the field to query. Can be dict or typed model instances
                (DenseVectorQuery, SparseVectorQuery, TextQuery, QueryStringQuery).
            include_fields: Fields to include in response documents. Use "*" for all
                fields, provide a list of field names for specific fields, or None
                to return document IDs only (default).
            filter: Metadata filter to narrow results using MongoDB query operators:
                $eq (equals), $ne (not equals), $gt/$gte/$lt/$lte (comparisons),
                $in/$nin (array membership), $and/$or (logical). Optional.
                Example: {"genre": {"$eq": "sci-fi"}, "year": {"$gte": 2020}}

        Returns:
            DocumentSearchResponse containing matching documents (sorted by relevance),
            namespace, and usage statistics.

        Raises:
            ValueError: If namespace is empty, top_k is out of range (1-10000),
                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:
            Dense vector search (using typed models):

            >>> from pinecone import Pinecone
            >>> from pinecone.models import DenseVectorQuery
            >>> pc = Pinecone(api_key="your-api-key")
            >>> index = pc.Index("example-index")
            >>> response = await index.documents.search(
            ...     namespace="movies",
            ...     top_k=10,
            ...     score_by=[DenseVectorQuery(field="_values", values=[0.1, 0.2, 0.3])],
            ...     include_fields=["text", "category"],
            ... )
            >>> for doc in response.matches:
            ...     print(f"{doc.id}: {doc.score}")

            Dense vector search (using dict):

            >>> response = await index.documents.search(
            ...     namespace="movies",
            ...     top_k=10,
            ...     score_by=[{"type": "dense_vector", "field": "_values", "values": [0.1, 0.2, 0.3]}],
            ...     include_fields=["text", "category"],
            ... )

            Sparse vector search:

            >>> from pinecone.models import SparseVectorQuery, SparseValues
            >>> response = await index.documents.search(
            ...     namespace="movies",
            ...     top_k=10,
            ...     score_by=[
            ...         SparseVectorQuery(
            ...             field="_sparse_values",
            ...             sparse_values=SparseValues(
            ...                 indices=[0, 5, 10, 100], values=[0.5, 0.3, 0.8, 0.2]
            ...             ),
            ...         )
            ...     ],
            ...     include_fields=["text"],
            ... )

            Full-text search:

            >>> from pinecone.models import TextQuery
            >>> response = await index.documents.search(
            ...     namespace="movies",
            ...     top_k=5,
            ...     score_by=[TextQuery(field="text", query="action movie with robots")],
            ...     include_fields=["text"],
            ... )

            Query string search:

            >>> from pinecone.models import QueryStringQuery
            >>> response = await index.documents.search(
            ...     namespace="movies",
            ...     top_k=10,
            ...     score_by=[QueryStringQuery(field="text", query="robots AND (action OR adventure)")],
            ...     include_fields=["text", "genre"],
            ... )

            With metadata filter:

            >>> response = await index.documents.search(
            ...     namespace="movies",
            ...     top_k=10,
            ...     score_by=[DenseVectorQuery(field="_values", values=[0.1, 0.2, 0.3])],
            ...     filter={"genre": {"$eq": "sci-fi"}},
            ...     include_fields=["text", "genre"],
            ... )

            Hybrid search (combining multiple query types):

            >>> from pinecone.models import DenseVectorQuery, TextQuery
            >>> response = await index.documents.search(
            ...     namespace="movies",
            ...     top_k=10,
            ...     score_by=[
            ...         DenseVectorQuery(field="_values", values=[0.1, 0.2, 0.3]),
            ...         TextQuery(field="text", query="robot action"),
            ...     ],
            ...     include_fields=["text", "genre"],
            ... )

        Note:
            Returns an empty list of matches when no documents match the query criteria.
            The namespace field will still be populated.

            Large top_k values (>1000) may impact response time. Use include_fields
            to limit data transfer when you don't need full documents.
        """
        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."
            )

        # Optimize: Check if we have any Structs before conversion
        has_structs = any(isinstance(q, Struct) for q in score_by)

        if has_structs:
            # Convert Structs to dicts
            normalized_score_by = [
                msgspec.to_builtins(q) if isinstance(q, Struct) else q for q in score_by
            ]
        else:
            # Fast path: all dicts, no conversion needed
            normalized_score_by = score_by

        body: dict[str, Any] = {
            "top_k": top_k,
            "score_by": normalized_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,
        )

        # Parse response - API returns "matches", "namespace", "usage"
        raw = orjson.loads(data)
        matches = [Document(doc) for doc in raw.get("matches", [])]
        namespace_value = raw["namespace"]
        usage_data = raw.get("usage")
        usage = Usage(read_units=usage_data.get("read_units")) if usage_data else None
        return DocumentSearchResponse(matches=matches, namespace=namespace_value, usage=usage)

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