"""
Compass parser client for document parsing operations.

This module provides the CompassParserClient class which handles document
parsing operations including file processing, metadata extraction, and
integration with various filesystems via fsspec.
"""

# Python imports
import json
import logging
from collections.abc import Callable, Iterable
from concurrent.futures import ThreadPoolExecutor
from datetime import timedelta
from types import TracebackType
from typing import Any

# 3rd party imports
import httpx

# 3rd party imports
from tenacity import retry, retry_if_exception, stop_after_attempt, wait_fixed

# Local imports
from cohere_compass import ProcessFileParameters
from cohere_compass.constants import (
    DEFAULT_COMPASS_PARSER_CLIENT_TIMEOUT,
    DEFAULT_MAX_RETRIES,
    DEFAULT_RETRY_WAIT,
)
from cohere_compass.exceptions import handle_httpx_exceptions
from cohere_compass.models import CompassDocument, ParserConfig
from cohere_compass.utils.fs import open_document, scan_folder
from cohere_compass.utils.iter import imap_parallel
from cohere_compass.utils.retry import is_retryable_compass_exception

Fn_or_Dict = dict[str, Any] | Callable[[CompassDocument], dict[str, Any]]


logger = logging.getLogger(__name__)


class CompassParserClient:
    """
    Client to interact with the CompassParser API.

    Handles document parsing operations with support for multiple file formats (PDF,
    DOCX, JSON, CSV, etc.) and integrates with various filesystems (local, S3, GCS) via
    fsspec. The client maintains default parser and metadata configurations that can be
    overridden per file.

    Attributes:
        parser_url: URL of the CompassParser API.
        parser_config: Default parser configuration for processing files.
        bearer_token: Optional authentication token.
        num_workers: Number of parallel workers for processing.
        timeout: Request timeout duration.

    """

    def __init__(
        self,
        *,
        parser_url: str,
        parser_config: ParserConfig = ParserConfig(),
        bearer_token: str | None = None,
        num_workers: int = 1,
        timeout: timedelta | None = None,
        httpx_client: httpx.Client | None = None,
    ):
        """
        Initialize the CompassParserClient.

        The parser_config is optional, and if not provided, the default configuration
        will be used. If provided, it will be used for all subsequent files processed by
        the client unless specific configs are passed when calling the process_file or
        process_files methods.

        :param parser_url: the URL of the CompassParser API
        :param parser_config: the parser configuration to use when processing files if
            no parser configuration is specified in the method calls (process_file or
            process_files)
        :param bearer_token (optional): The bearer token for authentication.
        :param num_workers (optional): The number of workers to use for processing
            files.
        :param timeout (optional): The timeout to use for the httpx client. Default is
            DEFAULT_COMPASS_PARSER_CLIENT_TIMEOUT.
        :param httpx_client (optional): 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.
        """
        self.parser_url = parser_url if not parser_url.endswith("/") else parser_url[:-1]
        self.parser_config = parser_config
        self.bearer_token = bearer_token
        self.thread_pool = ThreadPoolExecutor(num_workers)
        self.num_workers = num_workers
        self.timeout = (
            timeout
            if timeout is not None
            else (
                DEFAULT_COMPASS_PARSER_CLIENT_TIMEOUT
                if httpx_client is None
                else (
                    timedelta(seconds=httpx_client.timeout.read)
                    if httpx_client.timeout.read
                    else DEFAULT_COMPASS_PARSER_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

        logger.info(f"CompassParserClient initialized with parser_url: {self.parser_url}")

    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 process_folder(
        self,
        *,
        folder_path: str,
        allowed_extensions: list[str] | None = None,
        recursive: bool = False,
        parser_config: ParserConfig | None = None,
        custom_context: Fn_or_Dict | None = None,
    ):
        """
        Process all the files in the specified folder.

        The files are processed using the default parser configuration passed when
        creating the client. The method iterates over all the files in the folder and
        processes them using the process_file method. The resulting documents are
        returned as a list of CompassDocument objects.

        :param folder_path: the folder to process
        :param allowed_extensions: the list of allowed extensions to process
        :param recursive: whether to process the folder recursively
        :param parser_config: the parser configuration to use when processing files if
            no parser configuration is specified in the method calls (process_file or
            process_files)
        :param custom_context: Additional data to add to compass document. Fields will
            be filterable but not semantically searchable.  Can either be a dictionary
            or a callable that takes a CompassDocument and returns a dictionary.

        :return: the list of processed documents
        """
        filenames = scan_folder(
            folder_path=folder_path,
            allowed_extensions=allowed_extensions,
            recursive=recursive,
        )
        return self.process_files(
            filenames=filenames,
            parser_config=parser_config,
            custom_context=custom_context if custom_context else None,
        )

    def process_files(
        self,
        *,
        filenames: list[str],
        file_ids: list[str] | None = None,
        parser_config: ParserConfig | None = None,
        custom_context: Fn_or_Dict | None = None,
    ) -> Iterable[CompassDocument | tuple[str, Exception]]:
        """
        Process a list of files.

        If the parser config is not provided, then the default config passed by parameter
        when creating the client will be used. This makes the CompassParserClient
        stateful. That is, we can set the parser config only once when creating the
        parser client, and process all subsequent files without having to pass the
        config every time.

        All the documents passed as filenames and opened to obtain their bytes. Then,
        they are packed into a ProcessFilesParameters object that contains a list of
        ProcessFileParameters, each contain a file, its id, and the parser config.

        :param filenames: List of filenames to process
        :param file_ids: List of ids for the files
        :param parser_config: ParserConfig object (applies the same config to all docs)
        :param custom_context: Additional data to add to compass document. Fields will
            be filterable but not semantically searchable.  Can either be a dictionary
            or a callable that takes a CompassDocument and returns a dictionary.

        :return: List of processed documents
        """

        def process_file(i: int) -> list[CompassDocument] | tuple[str, Exception]:
            filename = filenames[i]
            try:
                return self.process_file(
                    filename=filename,
                    file_id=file_ids[i] if file_ids else None,
                    parser_config=parser_config,
                    custom_context=custom_context,
                )
            except Exception as e:
                return filename, e

        for results in imap_parallel(
            self.thread_pool,
            process_file,
            range(len(filenames)),
            max_parallelism=self.num_workers,
        ):
            if isinstance(results, list):
                yield from results
            else:
                yield results

    @staticmethod
    def _get_metadata(doc: CompassDocument, custom_context: Fn_or_Dict | None = None) -> dict[str, Any]:
        if custom_context is None:
            return {}
        elif callable(custom_context):
            return custom_context(doc)
        else:
            return custom_context

    def process_file(
        self,
        *,
        filename: str,
        file_id: str | None = None,
        content_type: str | None = None,
        parser_config: ParserConfig | None = None,
        custom_context: Fn_or_Dict | None = None,
        timeout: timedelta | None = None,
    ) -> list[CompassDocument]:
        """
        Process a file.

        The method takes in a file, its id, and the parser config. If the config is None,
        then it uses the default config passed by parameter when creating the client.
        This makes the CompassParserClient stateful for convenience, that is, one can
        pass in the parser config only once when creating the CompassParserClient, and
        process files without having to pass the config every time.

        :param filename: Filename to process.
        :param file_id: Id for the file.
        :param content_type: Content type of the file.
        :param parser_config: ParserConfig object with the config to use for parsing the
            file.
        :param custom_context: Additional data to add to compass document. Fields will
            be filterable but not semantically searchable.  Can either be a dictionary
            or a callable that takes a CompassDocument and returns a dictionary.
        :param timeout: Timeout in seconds for the process_file request. If None, uses
            the timeout set when creating the client.

        :return: List of resulting documents
        """
        doc = open_document(filename)
        if doc.errors:
            logger.error(f"Error opening document: {doc.errors}")
            return []

        return self._process_file_bytes(
            params=self._get_file_params(
                parser_config=parser_config,
                file_id=file_id,
                content_type=content_type,
            ),
            filename=filename,
            file_bytes=doc.filebytes,
            custom_context=custom_context,
            timeout=timeout,
        )

    def process_file_bytes(
        self,
        *,
        filename: str,
        file_bytes: bytes,
        file_id: str | None = None,
        content_type: str | None = None,
        parser_config: ParserConfig | None = None,
        custom_context: Fn_or_Dict | None = None,
        timeout: timedelta | None = None,
    ) -> list[CompassDocument]:
        """
        Process a file.

        The method takes in a file, its id, its byte array,
        and the parser/enrichment config.
        If the config is None, then it uses the default configs passed by parameter when
        creating the client.  This makes the CompassParserClient stateful for
        convenience, that is, one can pass in the parser config only once when creating
        the CompassParserClient, and process files without having to pass the config
        every time.

        :param filename: filename to process.
        :param file_bytes: byte content of the file
        :param file_id: ID for the file.
        :param content_type: Content type of the file.
        :param parser_config: ParserConfig object with the config to use for parsing the
            file.
        :param custom_context: Additional data to add to compass document. Fields will
            be filterable but not semantically searchable.  Can either be a dictionary
            or a callable that takes a CompassDocument and returns a dictionary.
        :param timeout: Timeout in seconds for the process_file request. If None, uses
            the timeout set when creating the client.

        :return: List of resulting documents
        """
        return self._process_file_bytes(
            params=self._get_file_params(
                parser_config=parser_config,
                file_id=file_id,
                content_type=content_type,
            ),
            filename=filename,
            file_bytes=file_bytes,
            custom_context=custom_context,
            timeout=timeout,
        )

    def _get_file_params(
        self,
        *,
        parser_config: ParserConfig | None = None,
        file_id: str | None = None,
        content_type: str | None = None,
    ):
        parser_config = parser_config or self.parser_config
        return ProcessFileParameters(
            parser_config=parser_config,
            doc_id=file_id,
            content_type=content_type,
        )

    @retry(
        stop=stop_after_attempt(DEFAULT_MAX_RETRIES),
        wait=wait_fixed(DEFAULT_RETRY_WAIT),
        retry=retry_if_exception(is_retryable_compass_exception),
        reraise=True,
    )
    def _process_file_bytes(
        self,
        *,
        params: ProcessFileParameters,
        filename: str,
        file_bytes: bytes,
        custom_context: Fn_or_Dict | None = None,
        timeout: timedelta | None = None,
    ) -> list[CompassDocument]:
        headers = None
        if self.bearer_token:
            headers = {"Authorization": f"Bearer {self.bearer_token}"}

        with handle_httpx_exceptions():
            res = self.httpx.post(
                url=f"{self.parser_url}/v1/process_file",
                data={"data": json.dumps(params.model_dump())},
                files={"file": (filename, file_bytes)},
                headers=headers,
                timeout=(timeout or self.timeout).total_seconds(),
            )
            res.raise_for_status()

        docs: list[CompassDocument] = []
        for doc in res.json()["docs"]:
            if doc.get("errors"):
                logger.error(f"Error processing file {filename}: {doc['errors']}")
            else:
                compass_doc = CompassDocument.adapt_doc_id_compass_doc(doc)
                additional_metadata = CompassParserClient._get_metadata(doc=compass_doc, custom_context=custom_context)
                compass_doc.content = {**compass_doc.content, **additional_metadata}
                docs.append(compass_doc)

        return docs
