"""Snowflake Cortex Search connector."""

from dataclasses import dataclass, field
import logging
import threading
import typing

from snowflake.connector import connect, errors, SnowflakeConnection
from snowflake.connector.network import ReauthenticationRequest
from snowflake.core import exceptions as core_exceptions, Root
from snowflake.core.cortex.search_service import QueryResponse
from tenacity import (
    retry,
    retry_if_exception_type,
    stop_after_attempt,
    wait_exponential,
)

from .exceptions import CortexSearchConfigError, map_snowflake_error, RETRYABLE_ERRORS
from .filters import validate_filter
from .utils import load_private_key

__all__ = ['CortexSearchConfig', 'CortexSearchClient']
logger = logging.getLogger(__name__)

DEFAULT_LIMIT = 20


@dataclass
class CortexSearchConfig:
    """Cortex Search config class."""

    SNOWFLAKE_ACCOUNT: str
    SNOWFLAKE_USER: str
    SNOWFLAKE_DATABASE: str
    SNOWFLAKE_SCHEMA: str
    SNOWFLAKE_CORTEX_SEARCH_SERVICE_NAME: str
    SNOWFLAKE_PRIVATE_KEY_PATH: str | None = field(default=None, repr=False)
    SNOWFLAKE_KEY_PASSPHRASE: str | None = field(default=None, repr=False)
    SNOWFLAKE_PRIVATE_KEY: str | None = field(default=None, repr=False)
    SNOWFLAKE_CORTEX_SEARCH_SCORING_CONFIG: dict | None = None

    SNOWFLAKE_LOGIN_TIMEOUT: int | None = None
    SNOWFLAKE_NETWORK_TIMEOUT: int | None = None

    RETRY_WAIT_MULTIPLIER: int = 1
    RETRY_MAX_ATTEMPTS: int = 3
    RETRY_WAIT_MIN: int = 1
    RETRY_WAIT_MAX: int = 30


class CortexSearchClient:
    """Snowflake Cortex Search client with key-pair JWT auth.

    Provides a thread-safe interface to query a Snowflake Cortex Search service
    using private key (JWT) authentication. The connection is established lazily
    on the first search call and reused across subsequent calls. Failed queries
    are automatically retried with exponential back-off for transient errors.

    Args:
        config (CortexSearchConfig): Configuration dataclass containing Snowflake
            credentials and retry settings. Either ``SNOWFLAKE_PRIVATE_KEY`` (key string)
            or ``SNOWFLAKE_PRIVATE_KEY_PATH`` (path to a PEM file) must
            be provided.
        logger_factory (Callable[[], logging.Logger] | None): Optional callable
            that returns a logger instance. Useful when the logger depends on
            request context. If ``None``, the module-level logger is used.

    Raises:
        CortexSearchConfigError: If neither ``SNOWFLAKE_PRIVATE_KEY`` nor
            ``SNOWFLAKE_PRIVATE_KEY_PATH`` is set, or if the key file cannot be
            loaded.

    Main methods:
        search: Submit search request.

    Example — basic setup::

        from payment.connectors.cortex_search import CortexSearchClient, CortexSearchConfig

        config = CortexSearchConfig(
            SNOWFLAKE_ACCOUNT='myorg-myaccount',
            SNOWFLAKE_USER='svc_user',
            SNOWFLAKE_DATABASE='MY_DB',
            SNOWFLAKE_SCHEMA='MY_SCHEMA',
            SNOWFLAKE_CORTEX_SEARCH_SERVICE_NAME='MY_SEARCH_SERVICE',
            SNOWFLAKE_PRIVATE_KEY_PATH='/path/to/rsa_key.p8',
            SNOWFLAKE_KEY_PASSPHRASE='secret',   # omit if key is not encrypted
        )
        client = CortexSearchClient(config, logger_factory=None)

    Example — querying with filters::

        from payment.connectors.cortex_search.filters import and_, eq_, gte_

        results = client.search(
            query='invoice payment overdue',
            columns=['contract_id', 'payment_date', 'amount'],
            query_filter=and_(
                eq_('status', 'OPEN'),
                gte_('payment_date', datetime.date(2025, 1, 1)),
            ),
            limit=10,
        )
        for row in results:
            print(row['contract_id'], row['amount'])

    Example — custom scoring config::

        results = client.search(
            query='tax withholding correction',
            columns=['document_id', 'title'],
            scoring_config={'semantic_weight': 0.8, 'bm25_weight': 0.2},
        )

    Example — closing the connection explicitly::

        client.close()   # releases the underlying Snowflake connection

    Thread safety:
        Connection initialisation is guarded by an internal ``threading.Lock``,
        making it safe to share a single ``CortexSearchClient`` instance across
        multiple threads.

    Retry behaviour:
        Transient errors (network timeouts, service unavailable, etc.) are retried
        up to ``config.RETRY_MAX_ATTEMPTS`` times with exponential back-off
        between ``config.RETRY_WAIT_MIN`` and ``config.RETRY_WAIT_MAX`` seconds.
        Authentication errors (``ForbiddenError``, ``TokenExpiredError``) trigger
        a connection reset before the exception is re-raised.
    """

    _config: CortexSearchConfig
    _logger_factory: typing.Callable[[], logging.Logger] | None
    _connection: SnowflakeConnection | None
    _private_key: bytes | None

    def __init__(
        self,
        config: CortexSearchConfig,
        logger_factory: typing.Callable[[], logging.Logger | None] | None,
    ):
        """
        Init Cortex Search client.
        Args:
            config: Cortex Search config.
            logger_factory: Callable to get logger in case it depends on some context.
        """
        self._config = config
        self._logger_factory = logger_factory
        self._connection = None
        self._service = None
        self._lock = threading.Lock()
        self._private_key = self._ensure_private_key()
        retry_decorator = retry(
            retry=retry_if_exception_type(RETRYABLE_ERRORS),
            wait=wait_exponential(
                multiplier=config.RETRY_WAIT_MULTIPLIER,
                min=config.RETRY_WAIT_MIN,
                max=config.RETRY_WAIT_MAX,
            ),
            stop=stop_after_attempt(config.RETRY_MAX_ATTEMPTS),
            reraise=True,
        )
        self.search = retry_decorator(self._search)

    def _ensure_private_key(self) -> bytes:
        """Ensure private key."""
        if (
            not self._config.SNOWFLAKE_PRIVATE_KEY
            and not self._config.SNOWFLAKE_PRIVATE_KEY_PATH
        ):
            raise CortexSearchConfigError(
                'Either SNOWFLAKE_PRIVATE_KEY or SNOWFLAKE_PRIVATE_KEY_PATH must be set'
            )

        try:
            private_key = load_private_key(
                self._config.SNOWFLAKE_PRIVATE_KEY,
                self._config.SNOWFLAKE_PRIVATE_KEY_PATH,
                self._config.SNOWFLAKE_KEY_PASSPHRASE,
            )
        except Exception:
            # ``from None`` to ensure there is no any chain or context that could
            # be exposed in logs or sentry
            raise CortexSearchConfigError('Error loading private key') from None

        return private_key

    @property
    def service(self):
        """Init service."""
        with self._lock:
            if self._connection is None:
                self._connection = connect(
                    account=self._config.SNOWFLAKE_ACCOUNT,
                    user=self._config.SNOWFLAKE_USER,
                    private_key=self._private_key,
                    authenticator='snowflake_jwt',
                    login_timeout=self._config.SNOWFLAKE_LOGIN_TIMEOUT,
                    network_timeout=self._config.SNOWFLAKE_NETWORK_TIMEOUT,
                )
            if self._service is None:
                self._service = (
                    Root(self._connection)
                    .databases[self._config.SNOWFLAKE_DATABASE]
                    .schemas[self._config.SNOWFLAKE_SCHEMA]
                    .cortex_search_services[
                        self._config.SNOWFLAKE_CORTEX_SEARCH_SERVICE_NAME
                    ]
                )
        return self._service

    @property
    def logger(self):
        return (self._logger_factory() or logger) if self._logger_factory else logger

    def _search(
        self,
        query: str,
        columns: list[str],
        query_filter: dict | None = None,
        scoring_config: dict | None = None,
        limit: int = DEFAULT_LIMIT,
    ) -> list[dict]:
        """Query Cortex Search service and return result rows.

        Args:
            query: Unstructured text query.
            columns: Columns to include in results.
            query_filter: Filter object built with and_/or_/eq_/gte_/lte_/not_.
            scoring_config: Scoring configuration.
            limit: Max number of results.

        Returns:
            List of result row dicts.
        """
        validate_filter(query_filter)

        try:
            response: QueryResponse = self.service.search(
                query=query,
                columns=columns,
                filter=query_filter,
                scoring_config=scoring_config
                or self._config.SNOWFLAKE_CORTEX_SEARCH_SCORING_CONFIG,
                limit=limit,
            )
        except (
            errors.ForbiddenError,
            errors.TokenExpiredError,
            core_exceptions.UnauthorizedError,
            core_exceptions.ForbiddenError,
            ReauthenticationRequest,
        ) as e:
            self.logger.warning('Cortex Search auth error, resetting connection: %s', e)
            # try to reinit connection / token in case of possible auth issues
            self.close()
            raise map_snowflake_error(e) from e
        except (errors.Error, core_exceptions.APIError) as e:
            raise map_snowflake_error(e) from e
        self.logger.debug('Cortex Search request_id=%s', response.request_id)
        return response.results

    def close(self) -> None:
        """Close the connection."""
        with self._lock:
            if self._connection is not None:
                try:
                    self._connection.close()
                except Exception as e:
                    self.logger.error('Error closing Cortex Search connection: %s', e)
                finally:
                    self._connection = None
            if self._service is not None:
                self._service = None