"""
Compass client for document indexing and search operations.

This module provides the CompassClient class which handles document indexing, searching,
and management operations with the Compass API. It includes support for batch
operations, retry logic, and comprehensive error handling.
"""

# Python imports
import base64
import logging
import os
import re
import threading
import uuid
from collections import deque
from collections.abc import Iterable
from dataclasses import dataclass
from datetime import timedelta
from statistics import mean
from types import TracebackType
from typing import Any, Literal

# 3rd party imports
import httpx
from joblib import Parallel, delayed  # type: ignore
from pydantic import BaseModel
from tenacity import (
    retry,
    retry_if_exception,
    stop_after_attempt,
    wait_fixed,
)

# Local imports
from cohere_compass import GroupAuthorizationInput
from cohere_compass.constants import (
    DEFAULT_COMPASS_CLIENT_TIMEOUT,
    DEFAULT_MAX_CHUNKS_PER_REQUEST,
    DEFAULT_MAX_ERROR_RATE,
    DEFAULT_MAX_RETRIES,
    DEFAULT_RETRY_WAIT,
    URL_SAFE_STRING_PATTERN,
)
from cohere_compass.exceptions import (
    CompassError,
    CompassInsertionError,
    CompassMaxErrorRateExceeded,
    handle_httpx_exceptions,
)
from cohere_compass.models import (
    CompassDocument,
    CompassSdkStage,
    DirectSearchInput,
    DirectSearchResponse,
    DirectSearchScrollInput,
    Document,
    ParseableDocument,
    PutDocumentsInput,
    SearchChunksResponse,
    SearchDocumentsResponse,
    SearchFilter,
    SearchInput,
    UploadDocumentsInput,
    UploadDocumentsResult,
)
from cohere_compass.models.config import IndexConfig
from cohere_compass.models.documents import (
    AssetPresignedUrlDetails,
    AssetPresignedUrlRequest,
    BulkUploadDocumentsStatus,
    BulkUploadStatusRequest,
    ContentTypeEnum,
    DocumentAttributes,
    GetAssetPresignedUrlsRequest,
    GetAssetPresignedUrlsResponse,
    ParseableDocumentConfig,
    ParsedDocumentResponse,
    PutDocumentsResponse,
    UploadDocumentsStatus,
)
from cohere_compass.models.indexes import IndexDetails, ListIndexesResponse
from cohere_compass.models.search import (
    GetDocumentResponse,
    RetrievedDocument,
    SortBy,
)
from cohere_compass.utils.documents import partition_documents
from cohere_compass.utils.retry import is_retryable_httpx_exception


@dataclass
class _SendRequestResult:
    """
    Result of a retryable HTTP request operation.

    This internal class encapsulates the result of HTTP requests made by the client.
    It's used for internal error handling and retry logic.

    Attributes:
        result: The response data if successful, otherwise None.
        content_type: The content type of the response, None if not available.

    """

    result: str | bytes | dict[str, Any] | list[dict[str, Any]] | None = None
    content_type: str | None = None


logger = logging.getLogger(__name__)


API_DEFINITIONS = {
    # General APIs
    "get_models": (
        "GET",
        "config/models",
    ),
    # Index APIs
    "create_index": (
        "PUT",
        "indexes/{index_name}",
    ),
    "list_indexes": (
        "GET",
        "indexes",
    ),
    "get_index_details": (
        "GET",
        "indexes/{index_name}",
    ),
    "delete_index": (
        "DELETE",
        "indexes/{index_name}",
    ),
    "refresh": (
        "POST",
        "indexes/{index_name}/_refresh",
    ),
    "update_group_authorization": (
        "POST",
        "indexes/{index_name}/group_authorization",
    ),
    # Document APIs
    "delete_document": (
        "DELETE",
        "indexes/{index_name}/documents/{document_id}",
    ),
    "get_document": (
        "GET",
        "indexes/{index_name}/documents/{document_id}",
    ),
    "put_documents": (
        "PUT",
        "indexes/{index_name}/documents",
    ),
    "get_document_asset": (
        "GET",
        "indexes/{index_name}/documents/{document_id}/assets/{asset_id}",
    ),
    "get_asset_presigned_urls": (
        "POST",
        "indexes/{index_name}/assets/_presigned_urls",
    ),
    "add_attributes": (
        "POST",
        "indexes/{index_name}/documents/{document_id}/_add_attributes",
    ),
    "upload_documents": (
        "POST",
        "indexes/{index_name}/documents/upload",
    ),
    "upload_documents_status": (
        "GET",
        "indexes/{index_name}/documents/upload/{upload_id}",
    ),
    "bulk_upload_documents_status": (
        "POST",
        "indexes/{index_name}/documents/uploads",
    ),
    "download_parsed_document": (
        "GET",
        "indexes/{index_name}/documents/upload/{upload_id}/_download",
    ),
    # Search APIs
    "search_documents": (
        "POST",
        "indexes/{index_name}/documents/_search",
    ),
    "search_chunks": (
        "POST",
        "indexes/{index_name}/documents/_search_chunks",
    ),
    "direct_search": (
        "POST",
        "indexes/{index_name}/_direct_search",
    ),
    "direct_search_scroll": (
        "POST",
        "indexes/{index_name}/_direct_search/scroll",
    ),
    # Task APIs
    "get_task": (
        "GET",
        "tasks/{task_id}",
    ),
    # Get visual element from asset
    "get_visual_element": (
        "GET",
        "indexes/{index_name}/documents/{document_id}/assets/{asset_id}",
    ),
}


class CompassClient:
    """
    Client for interacting with the Compass API.

    Provides methods for document indexing, searching, and management operations.
    Supports batch operations with configurable parallelism and retry logic.
    """

    def __init__(
        self,
        *,
        index_url: str,
        bearer_token: str | None = None,
        max_retries: int = DEFAULT_MAX_RETRIES,
        retry_wait: timedelta = DEFAULT_RETRY_WAIT,
        timeout: timedelta | None = None,
        httpx_client: httpx.Client | None = None,
    ):
        """
        Initialize the Compass client.

        :param index_url: The base URL for the Compass API.
        :param bearer_token: Optional bearer token for API authentication.
        :param max_retries: Maximum number of retries for failed requests.
        :param retry_wait: Time to wait between retries.
        :param timeout: Request timeout duration. If not specified, it defaults to
            DEFAULT_COMPASS_CLIENT_TIMEOUT.
        :param httpx_client: The httpx client to use for making requests. If not
            provided, a new httpx client will be created with the timeout set when
            creating the client.

        :raises: ValueError: If max_retries is negative or retry_wait is negative.
        """
        self.index_url = index_url if index_url.endswith("/") else f"{index_url}/"
        self.timeout = (
            timeout
            if timeout is not None
            else DEFAULT_COMPASS_CLIENT_TIMEOUT
            if httpx_client is None
            else timedelta(seconds=httpx_client.timeout.read)
            if httpx_client.timeout.read
            else DEFAULT_COMPASS_CLIENT_TIMEOUT
        )
        self.httpx = httpx_client or httpx.Client(timeout=self.timeout.total_seconds())
        self._own_httpx_client = httpx_client is None
        self._closed = False

        self.bearer_token = bearer_token

        if max_retries < 0:
            raise ValueError("default_max_retries must be a non-negative integer.")
        if retry_wait.total_seconds() < 0:
            raise ValueError("retry_wait must be a non-negative integer.")
        self.max_retries = max_retries
        self.retry_wait = retry_wait

    def close(self):
        """Close the httpx client if it was created by this instance."""
        if self._own_httpx_client and not self._closed:
            self.httpx.close()
            self._closed = True

    def __enter__(self):
        """For use by "with" statements."""
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_value: BaseException | None,
        traceback: TracebackType | None,
    ) -> None:
        """For use by "with" statements."""
        self.close()

    def get_models(
        self,
        max_retries: int | None = None,
        retry_wait: timedelta | None = None,
        timeout: timedelta | None = None,
    ) -> dict[str, list[str]]:
        """
        Get the models available in Compass.

        :param max_retries: Maximum number of retries for failed requests. If not
            provided, the default from the client will be used.
        :param retry_wait: Time to wait between retries. If not provided, the default
            from the client will be used.
        :param timeout: Request timeout duration. If not provided, the default from the
            client will be used.

        :return: Dictionary with model roles as keys ("dense", "rerank", "sparse") and
            lists of model versions as values.
        """
        result = self._send_request(
            api_name="get_models",
            max_retries=max_retries,
            retry_wait=retry_wait,
            timeout=timeout,
        )
        if not isinstance(result.result, dict):
            raise ValueError("Invalid response from Compass API")

        return result.result

    def create_index(
        self,
        *,
        index_name: str,
        index_config: IndexConfig | None = None,
        max_retries: int | None = None,
        retry_wait: timedelta | None = None,
        timeout: timedelta | None = None,
    ):
        """
        Create an index in Compass.

        :param index_name: The name of the index to create.
        :param index_config: Optional configuration for the index.
        :param max_retries: Maximum number of retries for failed requests. If not
            provided, the default from the client will be used.
        :param retry_wait: Time to wait between retries. If not provided, the default
            from the client will be used.
        :param timeout: Request timeout duration. If not provided, the default from the
            client will be used.

        :raises: ValueError: If index_name contains invalid characters.

        """
        if not re.match(URL_SAFE_STRING_PATTERN, index_name):
            raise ValueError(f"Invalid index name '{index_name}', please avoid special characters.")
        self._send_request(
            api_name="create_index",
            index_name=index_name,
            data=index_config,
            max_retries=max_retries,
            retry_wait=retry_wait,
            timeout=timeout,
        )

    def get_index_details(
        self,
        *,
        index_name: str,
        max_retries: int | None = None,
        retry_wait: timedelta | None = None,
        timeout: timedelta | None = None,
    ) -> IndexDetails:
        """
        Get the details of an index in Compass.

        :param index_name: The name of the index to query.
        :param max_retries: Maximum number of retries for failed requests. If not
            provided, the default from the client will be used.
        :param retry_wait: Time to wait between retries. If not provided, the default
            from the client will be used.
        :param timeout: Request timeout duration. If not provided, the default from the
            client will be used.


        Returns:
            IndexConfig object containing the index details.

        """
        result = self._send_request(
            api_name="get_index_details",
            index_name=index_name,
            max_retries=max_retries,
            retry_wait=retry_wait,
            timeout=timeout,
        )

        return IndexDetails.model_validate(result.result)

    def refresh_index(
        self,
        *,
        index_name: str,
        max_retries: int | None = None,
        retry_wait: timedelta | None = None,
        timeout: timedelta | None = None,
    ):
        """
        Refresh an index to make recent changes searchable.

        :param index_name: The name of the index to refresh.
        :param max_retries: Maximum number of retries for failed requests. If not
            provided, the default from the client will be used.
        :param retry_wait: Time to wait between retries. If not provided, the default
            from the client will be used.
        :param timeout: Request timeout duration. If not provided, the default from the
            client will be used.

        """
        self._send_request(
            api_name="refresh",
            index_name=index_name,
            max_retries=max_retries,
            retry_wait=retry_wait,
            timeout=timeout,
        )

    def delete_index(
        self,
        *,
        index_name: str,
        max_retries: int | None = None,
        retry_wait: timedelta | None = None,
        timeout: timedelta | None = None,
    ):
        """
        Delete an index from Compass.

        :param index_name: The name of the index to delete.
        :param max_retries: Maximum number of retries for failed requests. If not
            provided, the default from the client will be used.
        :param retry_wait: Time to wait between retries. If not provided, the default
            from the client will be used.
        :param timeout: Request timeout duration. If not provided, the default from the
            client will be used.

        """
        self._send_request(
            api_name="delete_index",
            index_name=index_name,
            max_retries=max_retries,
            retry_wait=retry_wait,
            timeout=timeout,
        )

    def delete_document(
        self,
        *,
        index_name: str,
        document_id: str,
        max_retries: int | None = None,
        retry_wait: timedelta | None = None,
        timeout: timedelta | None = None,
    ):
        """
        Delete a document from Compass.

        :param index_name: The name of the index containing the document.
        :param document_id: The ID of the document to delete.
        :param max_retries: Maximum number of retries for failed requests. If not
            provided, the default from the client will be used.
        :param retry_wait: Time to wait between retries. If not provided, the default
            from the client will be used.
        :param timeout: Request timeout duration. If not provided, the default from the
            client will be used.

        """
        self._send_request(
            api_name="delete_document",
            document_id=document_id,
            index_name=index_name,
            max_retries=max_retries,
            retry_wait=retry_wait,
            timeout=timeout,
        )

    def get_document(
        self,
        *,
        index_name: str,
        document_id: str,
        max_retries: int | None = None,
        retry_wait: timedelta | None = None,
        timeout: timedelta | None = None,
    ) -> RetrievedDocument:
        """
        Get a document from Compass.

        :param index_name: The name of the index containing the document.
        :param document_id: The ID of the document to retrieve.
        :param max_retries: Maximum number of retries for failed requests. If not
            provided, the default from the client will be used.
        :param retry_wait: Time to wait between retries. If not provided, the default
            from the client will be used.
        :param timeout: Request timeout duration. If not provided, the default from the
            client will be used.

        Returns:
            RetrievedDocument object containing the requested document.

        """
        result = self._send_request(
            api_name="get_document",
            document_id=document_id,
            index_name=index_name,
            max_retries=max_retries,
            retry_wait=retry_wait,
            timeout=timeout,
        )
        response = GetDocumentResponse.model_validate(result.result)
        return response.document

    def list_indexes(
        self,
        max_retries: int | None = None,
        retry_wait: timedelta | None = None,
        timeout: timedelta | None = None,
    ):
        """
        List all indexes in Compass.

        :param max_retries: Maximum number of retries for failed requests. If not
            provided, the default from the client will be used.
        :param retry_wait: Time to wait between retries. If not provided, the default
            from the client will be used.
        :param timeout: Request timeout duration. If not provided, the default from the
            client will be used.

        Returns:
            ListIndexesResponse containing all available indexes.

        """
        result = self._send_request(
            api_name="list_indexes",
            max_retries=max_retries,
            retry_wait=retry_wait,
            timeout=timeout,
        )

        return ListIndexesResponse.model_validate(result.result)

    def add_attributes(
        self,
        *,
        index_name: str,
        document_id: str,
        attributes: DocumentAttributes,
        max_retries: int | None = None,
        retry_wait: timedelta | None = None,
        timeout: timedelta | None = None,
    ):
        """
        Update the content field of an existing document with additional attributes.

        :param index_name: The name of the index containing the document.
        :param document_id: The ID of the document to modify.
        :param attributes: The attributes to add to the document.
        :param max_retries: Maximum number of retries for failed requests. If not
            provided, the default from the client will be used.
        :param retry_wait: Time to wait between retries. If not provided, the default
            from the client will be used.
        :param timeout: Request timeout duration. If not provided, the default from the
            client will be used.


        """
        self._send_request(
            api_name="add_attributes",
            document_id=document_id,
            data=attributes,
            index_name=index_name,
            max_retries=max_retries,
            retry_wait=retry_wait,
            timeout=timeout,
        )

    def insert_doc(
        self,
        *,
        index_name: str,
        doc: CompassDocument,
        authorized_groups: list[str] | None = None,
        merge_groups_on_conflict: bool = False,
        max_retries: int | None = None,
        retry_wait: timedelta | None = None,
        timeout: timedelta | None = None,
    ):
        """
        Insert a parsed document into an index in Compass.

        :param index_name: The name of the index.
        :param doc: The parsed compass document to insert.
        :param authorized_groups: Optional list of groups authorized to access this
            document.
        :param merge_groups_on_conflict: Whether to merge groups on conflict.
        :param max_retries: Maximum number of retries for failed requests. If not
            provided, the default from the client will be used.
        :param retry_wait: Time to wait between retries. If not provided, the default
            from the client will be used.
        :param timeout: Request timeout duration. If not provided, the default from the
            client will be used.

        """
        self.insert_docs(
            index_name=index_name,
            docs=[doc],
            authorized_groups=authorized_groups,
            merge_groups_on_conflict=merge_groups_on_conflict,
            max_retries=max_retries,
            retry_wait=retry_wait,
            timeout=timeout,
        )

    def upload_document(
        self,
        *,
        index_name: str,
        filename: str,
        filebytes: bytes,
        document_id: str,
        attributes: DocumentAttributes = DocumentAttributes(),
        config: ParseableDocumentConfig = ParseableDocumentConfig(),
        content_type: ContentTypeEnum | None = None,
        authorized_groups: list[str] | None = None,
        merge_groups_on_conflict: bool = False,
        max_retries: int | None = None,
        retry_wait: timedelta | None = None,
        timeout: timedelta | None = None,
    ) -> UploadDocumentsResult:
        """
        Parse and insert a document into an index in Compass.

        :param index_name: The name of the index.
        :param filename: The filename of the document.
        :param filebytes: The raw bytes of the document.
        :param content_type: optional content type of the document.
            Recommended to pass it otherwise auto-detected.
        :param document_id: The ID to assign to the document.
        :param attributes: Additional attributes to add to the document.
        :param config: Configuration for the document parsing.
        :param authorized_groups: When Document-level Security is enabled,
            Set authorized user groups to access the document.
            These groups should exist in RBAC. Default will make the document public.
        :param merge_groups_on_conflict: When Document-level Security is enabled,
            allows merging authorized groups if document already exists.
        :param max_retries: Maximum number of retries for failed requests. If not
            provided, the default from the client will be used.
        :param retry_wait: Time to wait between retries. If not provided, the default
            from the client will be used.
        :param timeout: Request timeout duration. If not provided, the default from the
            client will be used.

        Returns:
            UploadDocumentsResult containing the upload ID and status.

        """
        b64 = base64.b64encode(filebytes).decode("utf-8")
        doc = ParseableDocument(
            id=document_id,
            filename=filename,
            content_type=content_type,
            content_length_bytes=len(filebytes),
            content_encoded_bytes=b64,
            attributes=attributes,
            config=config,
        )

        result = self._send_request(
            api_name="upload_documents",
            data=UploadDocumentsInput(
                documents=[doc],
                authorized_groups=authorized_groups,
                merge_groups_on_conflict=merge_groups_on_conflict,
            ),
            index_name=index_name,
            max_retries=max_retries,
            retry_wait=retry_wait,
            timeout=timeout,
        )

        return UploadDocumentsResult.model_validate(result.result)

    def upload_document_status(
        self,
        *,
        index_name: str,
        upload_id: uuid.UUID,
        max_retries: int | None = None,
        retry_wait: timedelta | None = None,
        timeout: timedelta | None = None,
    ) -> list[UploadDocumentsStatus]:
        """
        Get status of document upload.

        :param index_name: The name of the index.
        :param upload_id: The upload ID returned when uploading the document.
        :param max_retries: Maximum number of retries for failed requests. If not
            provided, the default from the client will be used.
        :param retry_wait: Time to wait between retries. If not provided, the default
            from the client will be used.
        :param timeout: Request timeout duration. If not provided, the default from the
            client will be used.

        Returns:
            List of UploadDocumentsStatus objects containing upload status information.

        """
        result = self._send_request(
            api_name="upload_documents_status",
            index_name=index_name,
            upload_id=str(upload_id),
            max_retries=max_retries,
            retry_wait=retry_wait,
            timeout=timeout,
        )

        return [UploadDocumentsStatus(**r) for r in result.result]  # type: ignore

    def bulk_upload_document_status(
        self,
        *,
        index_name: str,
        upload_ids: list[uuid.UUID],
        max_retries: int | None = None,
        retry_wait: timedelta | None = None,
        timeout: timedelta | None = None,
    ) -> list[BulkUploadDocumentsStatus]:
        """
        Get status of multiple document uploads in a single request.

        :param index_name: The name of the index.
        :param upload_ids: List of upload IDs to fetch statuses for.
        :param max_retries: Maximum number of retries for failed requests. If not
            provided, the default from the client will be used.
        :param retry_wait: Time to wait between retries. If not provided, the default
            from the client will be used.
        :param timeout: Request timeout duration. If not provided, the default from the
            client will be used.

        Returns:
            List of BulkUploadDocumentsStatus objects, one per upload ID, preserving
            the order of the input upload_ids. Upload IDs with no matching files return
            an empty statuses list.

        """
        result = self._send_request(
            api_name="bulk_upload_documents_status",
            data=BulkUploadStatusRequest(upload_ids=upload_ids),
            index_name=index_name,
            max_retries=max_retries,
            retry_wait=retry_wait,
            timeout=timeout,
        )

        return [BulkUploadDocumentsStatus(**r) for r in result.result]  # type: ignore

    def download_parsed_document(
        self,
        *,
        index_name: str,
        upload_id: uuid.UUID,
        max_retries: int | None = None,
        retry_wait: timedelta | None = None,
        timeout: timedelta | None = None,
    ) -> list[ParsedDocumentResponse]:
        """
        Download the parsed document from Compass.

        :param index_name: The name of the index.
        :param upload_id: The upload ID returned when uploading the document.
        :param max_retries: Maximum number of retries for failed requests. If not
            provided, the default from the client will be used.
        :param retry_wait: Time to wait between retries. If not provided, the default
            from the client will be used.
        :param timeout: Request timeout duration. If not provided, the default from the
            client will be used.

        Returns:
            List of ParsedDocumentResponse objects containing the parsed documents.

        """
        result = self._send_request(
            api_name="download_parsed_document",
            index_name=index_name,
            upload_id=str(upload_id),
            max_retries=max_retries,
            retry_wait=retry_wait,
            timeout=timeout,
        )

        return [ParsedDocumentResponse.convert(data=r) for r in result.result]  # type: ignore

    def insert_docs(
        self,
        *,
        index_name: str,
        docs: Iterable[CompassDocument | tuple[str, Exception]],
        max_chunks_per_request: int = DEFAULT_MAX_CHUNKS_PER_REQUEST,
        max_error_rate: float = DEFAULT_MAX_ERROR_RATE,
        errors_sliding_window_size: int | None = 10,
        skip_first_n_docs: int = 0,
        num_jobs: int | None = None,
        authorized_groups: list[str] | None = None,
        merge_groups_on_conflict: bool = False,
        max_retries: int | None = None,
        retry_wait: timedelta | None = None,
        timeout: timedelta | None = None,
    ):
        """
        Insert multiple parsed documents into an index in Compass.

        :param index_name: The name of the index.
        :param docs: Iterator of parsed documents to insert.
        :param max_chunks_per_request: Maximum number of chunks to send in a single
            request.
        :param max_error_rate: Maximum error rate allowed before stopping insertion.
        :param errors_sliding_window_size: Size of the sliding window to track errors.
        :param skip_first_n_docs: Number of docs to skip indexing (useful for resuming
            failed insertions).
        :param num_jobs: Number of parallel jobs to use for insertion.
        :param authorized_groups: Groups authorized to access the documents (None makes
            documents public).
        :param merge_groups_on_conflict: Whether to merge groups when upserting
            documents with security.
        :param max_retries: Maximum number of retries for failed requests. If not
            provided, the default from the client will be used.
        :param retry_wait: Time to wait between retries. If not provided, the default
            from the client will be used.
        :param timeout: Request timeout duration. If not provided, the default from the
            client will be used.


        Returns:
            List of error details if any errors occurred, None otherwise.

        Raises:
            CompassMaxErrorRateExceeded: If error rate exceeds max_error_rate threshold.

        """

        def put_request(
            request_data: list[tuple[CompassDocument, Document]],
            previous_errors: list[dict[str, str]],
            num_doc: int,
        ) -> None:
            nonlocal num_succeeded, errors
            errors.extend(previous_errors)
            compass_docs: list[CompassDocument] = [compass_doc for compass_doc, _ in request_data]
            put_docs_input = PutDocumentsInput(
                documents=[input_doc for _, input_doc in request_data],
                authorized_groups=authorized_groups,
                merge_groups_on_conflict=merge_groups_on_conflict,
            )

            # It could be that all documents have errors, in which case we should not
            # send a request to the Compass Server. This is a common case when the
            # parsing of the documents fails.  In this case, only errors will appear in
            # the insertion_docs response
            if not request_data:
                return

            try:
                self._send_request(
                    api_name="put_documents",
                    data=put_docs_input,
                    index_name=index_name,
                    max_retries=max_retries,
                    retry_wait=retry_wait,
                    timeout=timeout,
                )
                num_succeeded += len(compass_docs)
            except CompassError as e:
                error = str(e)
                for doc in compass_docs:
                    filename = doc.metadata.filename
                    doc.errors.append({CompassSdkStage.Indexing: f"{filename}: {error}"})
                    errors.append({doc.metadata.document_id: f"{filename}: {error}"})

                # Keep track of the results of the last N API calls to calculate the
                # error rate If the error rate is higher than the threshold, stop the
                # insertion process
                error_window.append(error)

            error_rate = mean([1 if x else 0 for x in error_window]) if len(error_window) == error_window.maxlen else 0
            if error_rate > max_error_rate:
                raise CompassMaxErrorRateExceeded(
                    f"[Thread {threading.get_native_id()}] {error_rate * 100}% of "
                    f"insertions failed in the last {errors_sliding_window_size} API "
                    "calls. Stopping the insertion process."
                )

        error_window: deque[str | None] = deque(
            maxlen=errors_sliding_window_size
        )  # Keep track of the results of the last N API calls
        num_succeeded = 0
        errors: list[dict[str, str]] = []
        requests_iter = partition_documents(docs, max_chunks_per_request)

        try:
            num_jobs = num_jobs or os.cpu_count()
            Parallel(n_jobs=num_jobs, backend="threading")(
                delayed(put_request)(
                    request_data=request_block,
                    previous_errors=previous_errors,
                    num_doc=i,
                )
                for i, (request_block, previous_errors) in enumerate(requests_iter, 1)
                if i > skip_first_n_docs
            )
        except CompassMaxErrorRateExceeded as e:
            logger.error(e.message)

        if errors:
            raise CompassInsertionError(errors=errors)

    def _search(
        self,
        *,
        api_name: Literal["search_documents", "search_chunks"],
        index_name: str,
        query: str,
        top_k: int = 10,
        filters: list[SearchFilter] | None = None,
        rerank_model: str | None = None,
        max_retries: int | None = None,
        retry_wait: timedelta | None = None,
        timeout: timedelta | None = None,
    ):
        return self._send_request(
            api_name=api_name,
            index_name=index_name,
            data=SearchInput(query=query, top_k=top_k, filters=filters, rerank_model=rerank_model),
            max_retries=max_retries,
            retry_wait=retry_wait,
            timeout=timeout,
        )

    def search_documents(
        self,
        *,
        index_name: str,
        query: str,
        top_k: int = 10,
        filters: list[SearchFilter] | None = None,
        rerank_model: str | None = None,
        max_retries: int | None = None,
        retry_wait: timedelta | None = None,
        timeout: timedelta | None = None,
    ) -> SearchDocumentsResponse:
        """
        Search documents in an index.

        :param index_name: the name of the index
        :param query: the search query
        :param top_k: the number of documents to return
        :param filters: the search filters to apply
        :param rerank_model: the model to use for reranking the results
        :param max_retries: Maximum number of retries for failed requests. If not
            provided, the default from the client will be used.
        :param retry_wait: Time to wait between retries. If not provided, the default
            from the client will be used.
        :param timeout: Request timeout duration. If not provided, the default from the
            client will be used.

        Returns:
            SearchDocumentsResponse object containing the search results

        """
        result = self._search(
            api_name="search_documents",
            index_name=index_name,
            query=query,
            top_k=top_k,
            filters=filters,
            rerank_model=rerank_model,
            max_retries=max_retries,
            retry_wait=retry_wait,
            timeout=timeout,
        )

        return SearchDocumentsResponse.model_validate(result.result)

    def search_chunks(
        self,
        *,
        index_name: str,
        query: str,
        top_k: int = 10,
        filters: list[SearchFilter] | None = None,
        rerank_model: str | None = None,
        max_retries: int | None = None,
        retry_wait: timedelta | None = None,
        timeout: timedelta | None = None,
    ) -> SearchChunksResponse:
        """
        Search chunks in an index.

        :param index_name: the name of the index
        :param query: the search query
        :param top_k: the number of chunks to return
        :param filters: the search filters to apply
        :param rerank_model: the model to use for reranking the results
        :param max_retries: Maximum number of retries for failed requests. If not
            provided, the default from the client will be used.
        :param retry_wait: Time to wait between retries. If not provided, the default
            from the client will be used.
        :param timeout: Request timeout duration. If not provided, the default from the
            client will be used.

        Returns:
            SearchChunksResponse object containing the search results

        """
        result = self._search(
            api_name="search_chunks",
            index_name=index_name,
            query=query,
            top_k=top_k,
            filters=filters,
            rerank_model=rerank_model,
            max_retries=max_retries,
            retry_wait=retry_wait,
            timeout=timeout,
        )

        return SearchChunksResponse.model_validate(result.result)

    def get_document_asset(
        self,
        *,
        index_name: str,
        document_id: str,
        asset_id: str,
        max_retries: int | None = None,
        retry_wait: timedelta | None = None,
        timeout: timedelta | None = None,
    ) -> tuple[str | bytes | dict[str, Any], str]:
        """
        Get an asset from a document in Compass.

        :param index_name: the name of the index
        :param document_id: the id of the document
        :param asset_id: the id of the asset
        :param max_retries: Maximum number of retries for failed requests. If not
            provided, the default from the client will be used.
        :param retry_wait: Time to wait between retries. If not provided, the default
            from the client will be used.
        :param timeout: Request timeout duration. If not provided, the default from the
            client will be used.

        Returns:
            A tuple of the content and content type of the asset. The variable type of
            the content is either str, bytes, or dict[str, Any], depending on the asset
            type. For example, if the asset is an image, the content type will be bytes;
            if the asset is a markdown, the content type will be str; if the asset is a
            json, the content type will be dict[str, Any].

        Raises:
            CompassError: if the asset cannot be retrieved, either because it doesn't
                exist or the user doesn't have permission to access it.

        """
        result = self._send_request(
            api_name="get_document_asset",
            index_name=index_name,
            document_id=document_id,
            asset_id=asset_id,
            max_retries=max_retries,
            retry_wait=retry_wait,
            timeout=timeout,
        )

        return result.result, result.content_type  # type: ignore

    def get_asset_presigned_urls(
        self,
        *,
        index_name: str,
        assets: list[AssetPresignedUrlRequest],
        max_retries: int | None = None,
        retry_wait: timedelta | None = None,
        timeout: timedelta | None = None,
    ) -> list[AssetPresignedUrlDetails]:
        """
        Get presigned URLs for assets in documents.

        Retrieves presigned URLs for the specified assets. Each URL has a 1-hour TTL.
        This operation is all-or-nothing: if any requested asset pair is invalid or
        generation fails, the entire request fails.

        :param index_name: the name of the index
        :param assets: list of AssetPresignedUrlRequest objects containing document_id
            and asset_id pairs.
        :param max_retries: Maximum number of retries for failed requests. If not
            provided, the default from the client will be used.
        :param retry_wait: Time to wait between retries. If not provided, the default
            from the client will be used.
        :param timeout: Request timeout duration. If not provided, the default from the
            client will be used.

        Returns:
            A list of AssetPresignedUrlDetails objects in the same order as the input
            assets.

        Raises:
            CompassError: if any asset cannot be found or the user doesn't have READ
                permission on the target documents.
            CompassClientError: if the request body is malformed.

        """
        request_data = GetAssetPresignedUrlsRequest(assets=assets)

        result = self._send_request(
            api_name="get_asset_presigned_urls",
            index_name=index_name,
            data=request_data,
            max_retries=max_retries,
            retry_wait=retry_wait,
            timeout=timeout,
        )

        response = GetAssetPresignedUrlsResponse.model_validate(result.result)
        return response.asset_urls

    def update_group_authorization(
        self,
        *,
        index_name: str,
        group_auth_input: GroupAuthorizationInput,
        max_retries: int | None = None,
        retry_wait: timedelta | None = None,
        timeout: timedelta | None = None,
    ) -> PutDocumentsResponse:
        """
        Edit group authorization for an index.

        :param index_name: the name of the index
        :param group_auth_input: the group authorization input
        :param max_retries: Maximum number of retries for failed requests. If not
            provided, the default from the client will be used.
        :param retry_wait: Time to wait between retries. If not provided, the default
            from the client will be used.
        :param timeout: Request timeout duration. If not provided, the default from the
            client will be used.

        """
        result = self._send_request(
            api_name="update_group_authorization",
            index_name=index_name,
            data=group_auth_input,
            max_retries=max_retries,
            retry_wait=retry_wait,
            timeout=timeout,
        )
        return PutDocumentsResponse.model_validate(result.result)

    def direct_search(
        self,
        *,
        index_name: str,
        query: dict[str, Any],
        sort_by: list[SortBy] | None = None,
        size: int = 100,
        scroll: str | None = None,
        max_retries: int | None = None,
        retry_wait: timedelta | None = None,
        timeout: timedelta | None = None,
    ) -> DirectSearchResponse:
        """
        Perform a direct search query against the Compass API.

        :param index_name: the name of the index
        :param query: the direct search query (e.g. {"match_all": {}})
        :param sort_by: the sort by criteria
        :param size: the number of results to return
        :param scroll: the scroll duration (e.g. "1m" for 1 minute)
        :param max_retries: Maximum number of retries for failed requests. If not
            provided, the default from the client will be used.
        :param retry_wait: Time to wait between retries. If not provided, the default
            from the client will be used.
        :param timeout: Request timeout duration. If not provided, the default from the
            client will be used.

        :return: the direct search results
        :raises CompassError: if the search fails

        """
        data = DirectSearchInput(query=query, size=size, sort_by=sort_by, scroll=scroll)

        result = self._send_request(
            api_name="direct_search",
            index_name=index_name,
            data=data,
            max_retries=max_retries,
            retry_wait=retry_wait,
            timeout=timeout,
        )

        return DirectSearchResponse.model_validate(result.result)

    def direct_search_scroll(
        self,
        *,
        index_name: str,
        scroll_id: str,
        scroll: str = "1m",
        max_retries: int | None = None,
        retry_wait: timedelta | None = None,
        timeout: timedelta | None = None,
    ) -> DirectSearchResponse:
        """
        Continue a search using a scroll ID from a previous direct_search call.

        :param scroll_id: the scroll ID from a previous direct_search call
        :param index_name: the name of the index same as used in direct_search
        :param scroll: the scroll duration (e.g. "1m" for 1 minute)
        :param max_retries: Maximum number of retries for failed requests. If not
            provided, the default from the client will be used.
        :param retry_wait: Time to wait between retries. If not provided, the default
            from the client will be used.
        :param timeout: Request timeout duration. If not provided, the default from the
            client will be used.

        Returns:
            the next batch of search results

        Raises:
            CompassError: if the scroll search fails

        """
        data = DirectSearchScrollInput(scroll_id=scroll_id, scroll=scroll)
        result = self._send_request(
            api_name="direct_search_scroll",
            index_name=index_name,
            data=data,
            max_retries=max_retries,
            retry_wait=retry_wait,
            timeout=timeout,
        )

        return DirectSearchResponse.model_validate(result.result)

    def get_task(
        self,
        *,
        task_id: str,
        max_retries: int | None = None,
        retry_wait: timedelta | None = None,
        timeout: timedelta | None = None,
    ) -> dict[str, Any]:
        """
        Get a task by its ID.

        :param task_id: The ID of the task to retrieve.
        :param max_retries: Maximum number of retries for failed requests. If not
            provided, the default from the client will be used.
        :param retry_wait: Time to wait between retries. If not provided, the default
            from the client will be used.
        :param timeout: Request timeout duration. If not provided, the default from the
            client will be used.

        Returns:
            Dictionary containing the task details.

        """
        result = self._send_request(
            api_name="get_task",
            task_id=task_id,
            max_retries=max_retries,
            retry_wait=retry_wait,
            timeout=timeout,
        )
        if not isinstance(result.result, dict):
            raise ValueError("Invalid response from Compass API")

        return result.result

    def get_visual_element(
        self,
        *,
        index_name: str,
        document_id: str,
        asset_id: str,
        x0: int,
        y0: int,
        x1: int,
        y1: int,
        max_retries: int | None = None,
        retry_wait: timedelta | None = None,
        timeout: timedelta | None = None,
    ) -> tuple[str | bytes | dict[str, Any], str]:
        """
        Get a visual element from an asset in Compass by its bounding box coordinates.

        :param index_name: the name of the index
        :param document_id: the id of the document
        :param asset_id: the id of the asset
        :param x0: the x coordinate of the top left corner of the bounding box
        :param y0: the y coordinate of the top left corner of the bounding box
        :param x1: the x coordinate of the bottom right corner of the bounding box
        :param y1: the y coordinate of the bottom right corner of the bounding box
        :param max_retries: Maximum number of retries for failed requests. If not
            provided, the default from the client will be used.
        :param retry_wait: Time to wait between retries. If not provided, the default
            from the client will be used.
        :param timeout: Request timeout duration. If not provided, the default from the
            client will be used.

        Returns:
            tuple[bytes | str ]: bytes of image and content_type as string

        Raises:
            CompassError: if the visual element cannot be retrieved, either because it
                doesn't exist or the user doesn't have permission to access it.

        """
        result = self._send_request(
            api_name="get_visual_element",
            index_name=index_name,
            document_id=document_id,
            asset_id=asset_id,
            query_params={"x0": x0, "y0": y0, "x1": x1, "y1": y1},
            max_retries=max_retries,
            retry_wait=retry_wait,
            timeout=timeout,
        )

        return result.result, result.content_type  # type: ignore

    def _send_http_request(
        self,
        http_method: str,
        target_path: str,
        data: BaseModel | None = None,
        timeout: timedelta | None = None,
        query_params: dict[str, Any] | None = None,
    ):
        timeout = timeout or self.timeout

        data_dict = data.model_dump(mode="json", exclude_none=True) if data else None

        headers = None
        if self.bearer_token:
            headers = {"Authorization": f"Bearer {self.bearer_token}"}

        if http_method == "GET":
            response = self.httpx.get(
                target_path,
                headers=headers,
                params=query_params,
                timeout=timeout.total_seconds(),
            )
        elif http_method == "POST":
            response = self.httpx.post(
                target_path,
                json=data_dict,
                headers=headers,
                timeout=timeout.total_seconds(),
            )
        elif http_method == "PUT":
            response = self.httpx.put(
                target_path,
                json=data_dict,
                headers=headers,
                timeout=timeout.total_seconds(),
            )
        elif http_method == "DELETE":
            response = self.httpx.delete(
                target_path,
                headers=headers,
                timeout=timeout.total_seconds(),
            )
        else:
            raise RuntimeError(f"Unsupported HTTP method: {http_method}")

        response.raise_for_status()

        content_type = response.headers.get("content-type")
        if content_type in ("image/jpeg", "image/png"):
            # To handle response from get_document_asset() when the asset
            # is an image.
            result = response.content
        elif content_type == "text/markdown":
            # To handle response from get_document_asset() when the asset
            # is a markdown.
            result = response.text
        else:
            # To handle response from other APIs.
            result = response.json() if response.text else None
        return _SendRequestResult(
            result=result,
            content_type=content_type,
        )

    def _send_request(
        self,
        api_name: str,
        data: BaseModel | None = None,
        max_retries: int | None = None,
        retry_wait: timedelta | None = None,
        timeout: timedelta | None = None,
        query_params: dict[str, Any] | None = None,
        **url_params: str,
    ) -> _SendRequestResult:
        """
        Send a request to the Compass API.

        :param function: the function to call
        :param index_name: the name of the index
        :param data: the data to send
        :param max_retries: Maximum number of retries for failed requests. If not
            provided, the default from the client will be used.
        :param retry_wait: Time to wait between retries. If not provided, the default
            from the client will be used.
        :param timeout: Request timeout duration. If not provided, the default from the
            client will be used.
        :param query_params: Optional dictionary of query parameters to append to the URL.

        :return: An error message if the request failed, otherwise None.
        """
        max_retries = max_retries or self.max_retries
        retry_wait = retry_wait or self.retry_wait
        timeout = timeout or self.timeout

        if api_name not in API_DEFINITIONS:
            raise CompassError(f"API name '{api_name}' is not defined in the API definitions.")
        http_method, api_path = API_DEFINITIONS[api_name]

        target_path = f"{self.index_url}v1/{api_path}"
        target_path = target_path.format(**url_params)

        @retry(
            stop=stop_after_attempt(max_retries),
            wait=wait_fixed(retry_wait),
            reraise=True,
            retry=retry_if_exception(is_retryable_httpx_exception),
        )
        def _send_request_with_retry() -> _SendRequestResult:
            return self._send_http_request(
                http_method=http_method,
                target_path=target_path,
                data=data,
                timeout=timeout,
                query_params=query_params,
            )

        with handle_httpx_exceptions():
            return _send_request_with_retry()
