from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Any, Optional

from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.utils import get_from_env, guard_import

if TYPE_CHECKING:
    from whylogs.api.logger.logger import Logger

diagnostic_logger = logging.getLogger(__name__)


def import_langkit(
    sentiment: bool = False,
    toxicity: bool = False,
    themes: bool = False,
) -> Any:
    """Import the langkit python package and raise an error if it is not installed.

    Args:
        sentiment: Whether to import the langkit.sentiment module. Defaults to False.
        toxicity: Whether to import the langkit.toxicity module. Defaults to False.
        themes: Whether to import the langkit.themes module. Defaults to False.

    Returns:
        The imported langkit module.
    """
    langkit = guard_import("langkit")
    guard_import("langkit.regexes")
    guard_import("langkit.textstat")
    if sentiment:
        guard_import("langkit.sentiment")
    if toxicity:
        guard_import("langkit.toxicity")
    if themes:
        guard_import("langkit.themes")
    return langkit


class WhyLabsCallbackHandler(BaseCallbackHandler):
    """
    Callback Handler for logging to WhyLabs. This callback handler utilizes
    `langkit` to extract features from the prompts & responses when interacting with
    an LLM. These features can be used to guardrail, evaluate, and observe interactions
    over time to detect issues relating to hallucinations, prompt engineering,
    or output validation. LangKit is an LLM monitoring toolkit developed by WhyLabs.

    Here are some examples of what can be monitored with LangKit:
    * Text Quality
      - readability score
      - complexity and grade scores
    * Text Relevance
      - Similarity scores between prompt/responses
      - Similarity scores against user-defined themes
      - Topic classification
    * Security and Privacy
      - patterns - count of strings matching a user-defined regex pattern group
      - jailbreaks - similarity scores with respect to known jailbreak attempts
      - prompt injection - similarity scores with respect to known prompt attacks
      - refusals - similarity scores with respect to known LLM refusal responses
    * Sentiment and Toxicity
      - sentiment analysis
      - toxicity analysis

    For more information, see https://docs.whylabs.ai/docs/language-model-monitoring
    or check out the LangKit repo here: https://github.com/whylabs/langkit

    ---
    Args:
        api_key (Optional[str]): WhyLabs API key. Optional because the preferred
            way to specify the API key is with environment variable
            WHYLABS_API_KEY.
        org_id (Optional[str]): WhyLabs organization id to write profiles to.
            Optional because the preferred way to specify the organization id is
            with environment variable WHYLABS_DEFAULT_ORG_ID.
        dataset_id (Optional[str]): WhyLabs dataset id to write profiles to.
            Optional because the preferred way to specify the dataset id is
            with environment variable WHYLABS_DEFAULT_DATASET_ID.
        sentiment (bool): Whether to enable sentiment analysis. Defaults to False.
        toxicity (bool): Whether to enable toxicity analysis. Defaults to False.
        themes (bool): Whether to enable theme analysis. Defaults to False.
    """

    def __init__(self, logger: Logger, handler: Any):
        """Initiate the rolling logger."""
        super().__init__()
        if hasattr(handler, "init"):
            handler.init(self)
        if hasattr(handler, "_get_callbacks"):
            self._callbacks = handler._get_callbacks()
        else:
            self._callbacks = dict()
            diagnostic_logger.warning("initialized handler without callbacks.")
        self._logger = logger

    def flush(self) -> None:
        """Explicitly write current profile if using a rolling logger."""
        if self._logger and hasattr(self._logger, "_do_rollover"):
            self._logger._do_rollover()
            diagnostic_logger.info("Flushing WhyLabs logger, writing profile...")

    def close(self) -> None:
        """Close any loggers to allow writing out of any profiles before exiting."""
        if self._logger and hasattr(self._logger, "close"):
            self._logger.close()
            diagnostic_logger.info("Closing WhyLabs logger, see you next time!")

    def __enter__(self) -> WhyLabsCallbackHandler:
        return self

    def __exit__(
        self, exception_type: Any, exception_value: Any, traceback: Any
    ) -> None:
        self.close()

    @classmethod
    def from_params(
        cls,
        *,
        api_key: Optional[str] = None,
        org_id: Optional[str] = None,
        dataset_id: Optional[str] = None,
        sentiment: bool = False,
        toxicity: bool = False,
        themes: bool = False,
        logger: Optional[Logger] = None,
    ) -> WhyLabsCallbackHandler:
        """Instantiate whylogs Logger from params.

        Args:
            api_key (Optional[str]): WhyLabs API key. Optional because the preferred
                way to specify the API key is with environment variable
                WHYLABS_API_KEY.
            org_id (Optional[str]): WhyLabs organization id to write profiles to.
                If not set must be specified in environment variable
                WHYLABS_DEFAULT_ORG_ID.
            dataset_id (Optional[str]): The model or dataset this callback is gathering
                telemetry for. If not set must be specified in environment variable
                WHYLABS_DEFAULT_DATASET_ID.
            sentiment (bool): If True will initialize a model to perform
                sentiment analysis compound score. Defaults to False and will not gather
                this metric.
            toxicity (bool): If True will initialize a model to score
                toxicity. Defaults to False and will not gather this metric.
            themes (bool): If True will initialize a model to calculate
                distance to configured themes. Defaults to None and will not gather this
                metric.
            logger (Optional[Logger]): If specified will bind the configured logger as
                the telemetry gathering agent. Defaults to LangKit schema with periodic
                WhyLabs writer.
        """
        # langkit library will import necessary whylogs libraries
        import_langkit(sentiment=sentiment, toxicity=toxicity, themes=themes)

        why = guard_import("whylogs")
        get_callback_instance = guard_import(
            "langkit.callback_handler"
        ).get_callback_instance
        WhyLabsWriter = guard_import("whylogs.api.writer.whylabs").WhyLabsWriter
        udf_schema = guard_import("whylogs.experimental.core.udf_schema").udf_schema

        if logger is None:
            api_key = api_key or get_from_env("api_key", "WHYLABS_API_KEY")
            org_id = org_id or get_from_env("org_id", "WHYLABS_DEFAULT_ORG_ID")
            dataset_id = dataset_id or get_from_env(
                "dataset_id", "WHYLABS_DEFAULT_DATASET_ID"
            )
            whylabs_writer = WhyLabsWriter(
                api_key=api_key, org_id=org_id, dataset_id=dataset_id
            )

            whylabs_logger = why.logger(
                mode="rolling", interval=5, when="M", schema=udf_schema()
            )

            whylabs_logger.append_writer(writer=whylabs_writer)
        else:
            diagnostic_logger.info("Using passed in whylogs logger {logger}")
            whylabs_logger = logger

        callback_handler_cls = get_callback_instance(logger=whylabs_logger, impl=cls)
        diagnostic_logger.info(
            "Started whylogs Logger with WhyLabsWriter and initialized LangKit. 📝"
        )
        return callback_handler_cls