import datetime
import json
import logging
import os
import uuid
from typing import Any, Dict, List, Literal, Optional, Union

import certifi
from beekeeper.core.monitors import PromptMonitor
from beekeeper.core.monitors.types import PayloadRecord
from beekeeper.core.prompts.utils import extract_template_vars
from beekeeper.monitors.watsonx.instrumentation import suppress_output
from deprecated import deprecated
from pydantic.v1 import BaseModel

os.environ["REQUESTS_CA_BUNDLE"] = certifi.where()
logging.getLogger("ibm_watsonx_ai.client").setLevel(logging.ERROR)
logging.getLogger("ibm_watsonx_ai.wml_resource").setLevel(logging.ERROR)

REGIONS_URL = {
    "us-south": {
        "wml": "https://us-south.ml.cloud.ibm.com",
        "wos": "https://api.aiopenscale.cloud.ibm.com",
        "factsheet": None,
    },
    "eu-de": {
        "wml": "https://eu-de.ml.cloud.ibm.com",
        "wos": "https://eu-de.api.aiopenscale.cloud.ibm.com",
        "factsheet": "frankfurt",
    },
    "au-syd": {
        "wml": "https://au-syd.ml.cloud.ibm.com",
        "wos": "https://au-syd.api.aiopenscale.cloud.ibm.com",
        "factsheet": "sydney",
    },
}


def _filter_dict(original_dict: Dict, optional_keys: List, required_keys: List = []):
    """
    Filters a dictionary to keep only the specified keys and checks for required keys.

    Args:
        original_dict (Dict): The original dictionary.
        optional_keys (list): A list of keys to retain.
        required_keys (list, optional): A list of keys that must be present in the dictionary. Defaults to None.
    """
    # Ensure all required keys are in the source dict
    missing_keys = [key for key in required_keys if key not in original_dict]
    if missing_keys:
        raise KeyError(f"Missing required parameter: {missing_keys}")

    all_keys_to_keep = set(required_keys + optional_keys)

    # Create a new dictionary with only the key-value pairs where the key is in 'keys' and value is not None
    return {
        key: original_dict[key]
        for key in all_keys_to_keep
        if key in original_dict and original_dict[key] is not None
    }


def _convert_payload_format(
    records: List[Dict],
    feature_fields: List[str],
) -> List[Dict]:
    payload_data = []
    response_fields = ["generated_text", "input_token_count", "generated_token_count"]

    for record in records:
        request = {"parameters": {"template_variables": {}}}
        results = {}

        request["parameters"]["template_variables"] = {
            field: str(record.get(field, "")) for field in feature_fields
        }

        results = {
            field: record.get(field) for field in response_fields if record.get(field)
        }

        pl_record = {
            "request": request,
            "response": {"results": [results]},
            "scoring_id": str(uuid.uuid4()),
        }

        if "response_time" in record:
            pl_record["response_time"] = record["response_time"]

        payload_data.append(pl_record)

    return payload_data


# ===== Credentials Classes =====
class CloudPakforDataCredentials(BaseModel):
    """
    Encapsulates the credentials required for IBM Cloud Pak for Data.

    Attributes:
        url (str): The host URL of the Cloud Pak for Data environment.
        api_key (str, optional): The API key for the environment, if IAM is enabled.
        username (str, optional): The username for the environment.
        password (str, optional): The password for the environment.
        bedrock_url (str, optional): The Bedrock URL. Required only when IAM integration is enabled on CP4D 4.0.x clusters.
        instance_id (str, optional): The instance ID.
        version (str, optional): The version of Cloud Pak for Data.
        disable_ssl_verification (bool, optional): Indicates whether to disable SSL certificate verification.
            Defaults to `True`.
    """

    url: str
    api_key: Optional[str] = None
    username: Optional[str] = None
    password: Optional[str] = None
    bedrock_url: Optional[str] = None
    instance_id: Optional[Literal["icp", "openshift"]] = None
    version: Optional[str] = None
    disable_ssl_verification: bool = True

    def __init__(
        self,
        url: str,
        api_key: Optional[str] = None,
        username: Optional[str] = None,
        password: Optional[str] = None,
        bedrock_url: Optional[str] = None,
        instance_id: Optional[Literal["icp", "openshift"]] = None,
        version: Optional[str] = None,
        disable_ssl_verification: bool = True,
    ) -> None:
        super().__init__(
            url=url,
            api_key=api_key,
            username=username,
            password=password,
            bedrock_url=bedrock_url,
            instance_id=instance_id,
            version=version,
            disable_ssl_verification=disable_ssl_verification,
        )

    def to_dict(self) -> Dict[str, Any]:
        cpd_creds = dict([(k, v) for k, v in self.__dict__.items()])  # noqa: C404

        if "instance_id" in cpd_creds and self.instance_id.lower() not in [
            "icp",
            "openshift",
        ]:
            cpd_creds.pop("instance_id")

        return cpd_creds


class IntegratedSystemCredentials(BaseModel):
    """
    Encapsulates the credentials for an Integrated System based on the authentication type.

    Depending on the `auth_type`, only a subset of the properties is required.

    Attributes:
        auth_type (str): The type of authentication. Currently supports "basic" and "bearer".
        username (str, optional): The username for Basic Authentication.
        password (str, optional): The password for Basic Authentication.
        token_url (str, optional): The URL of the authentication endpoint used to request a Bearer token.
        token_method (str, optional): The HTTP method (e.g., "POST", "GET") used to request the Bearer token.
            Defaults to "POST".
        token_headers (Dict, optional): Optional headers to include when requesting the Bearer token.
            Defaults to `None`.
        token_payload (str | dict, optional): The body or payload to send when requesting the Bearer token.
            Can be a string (e.g., raw JSON). Defaults to `None`.
    """

    auth_type: Literal["basic", "bearer"]
    username: Optional[str]  # basic
    password: Optional[str]  # basic
    token_url: Optional[str]  # bearer
    token_method: Optional[str] = "POST"  # bearer
    token_headers: Optional[Dict] = {}  # bearer
    token_payload: Optional[Union[str, Dict]] = None  # bearer

    def __init__(
        self,
        auth_type: Literal["basic", "bearer"],
        username: str = None,
        password: str = None,
        token_url: str = None,
        token_method: str = "POST",
        token_headers: Dict = {},
        token_payload: Union[str, Dict] = None,
    ) -> None:
        if auth_type == "basic":
            if not username or not password:
                raise ValueError(
                    "`username` and `password` are required for auth_type = 'basic'.",
                )
        elif auth_type == "bearer":
            if not token_url:
                raise ValueError(
                    "`token_url` are required for auth_type = 'bearer'.",
                )

        super().__init__(
            auth_type=auth_type,
            username=username,
            password=password,
            token_url=token_url,
            token_method=token_method,
            token_headers=token_headers,
            token_payload=token_payload,
        )

    def to_dict(self) -> Dict:
        integrated_system_creds = {"auth_type": self.auth_type}

        if self.auth_type == "basic":
            integrated_system_creds["username"] = self.username
            integrated_system_creds["password"] = self.password
        elif self.auth_type == "bearer":
            integrated_system_creds["token_info"] = {
                "url": self.token_url,
                "method": self.token_method,
                "headers": self.token_headers,
                "payload": self.token_payload,
            }

        return integrated_system_creds


# ===== Monitor Classes =====
class WatsonxExternalPromptMonitor(PromptMonitor):
    """
    Provides functionality to interact with IBM watsonx.governance for monitoring prompts executed on external LLMs.

    Note:
        One of the following parameters is required to create a prompt monitor:
        `project_id` or `space_id`, but not both.

    Attributes:
        api_key (str): The API key for IBM watsonx.governance.
        space_id (str, optional): The space ID in watsonx.governance.
        project_id (str, optional): The project ID in watsonx.governance.
        region (str, optional): The region where watsonx.governance is hosted when using IBM Cloud.
            Defaults to `us-south`.
        cpd_creds (CloudPakforDataCredentials, optional): The Cloud Pak for Data environment credentials.
        subscription_id (str, optional): The subscription ID associated with the records being logged.

    Example:
        ```python
        from beekeeper.monitors.watsonx import (
            WatsonxExternalPromptMonitor,
            CloudPakforDataCredentials,
        )

        # watsonx.governance (IBM Cloud)
        wxgov_client = WatsonxExternalPromptMonitor(
            api_key="API_KEY", space_id="SPACE_ID"
        )

        # watsonx.governance (CP4D)
        cpd_creds = CloudPakforDataCredentials(
            url="CPD_URL",
            username="USERNAME",
            password="PASSWORD",
            version="5.0",
            instance_id="openshift",
        )

        wxgov_client = WatsonxExternalPromptMonitor(
            space_id="SPACE_ID", cpd_creds=cpd_creds
        )
        ```
    """

    def __init__(
        self,
        api_key: str = None,
        space_id: str = None,
        project_id: str = None,
        region: Literal["us-south", "eu-de", "au-syd"] = "us-south",
        cpd_creds: CloudPakforDataCredentials | Dict = None,
        subscription_id: str = None,
        **kwargs,
    ) -> None:
        import ibm_aigov_facts_client  # noqa: F401
        import ibm_cloud_sdk_core.authenticators  # noqa: F401
        import ibm_watson_openscale  # noqa: F401
        import ibm_watsonx_ai  # noqa: F401

        super().__init__(**kwargs)

        self.space_id = space_id
        self.project_id = project_id
        self.region = region
        self.subscription_id = subscription_id
        self._api_key = api_key
        self._wos_client = None

        self._container_id = space_id if space_id else project_id
        self._container_type = "space" if space_id else "project"
        self._deployment_stage = "production" if space_id else "development"

        if cpd_creds:
            self._wos_cpd_creds = _filter_dict(
                cpd_creds.to_dict(),
                ["username", "password", "api_key", "disable_ssl_verification"],
                ["url"],
            )
            self._fact_cpd_creds = _filter_dict(
                cpd_creds.to_dict(),
                ["username", "password", "api_key", "bedrock_url"],
                ["url"],
            )
            self._fact_cpd_creds["service_url"] = self._fact_cpd_creds.pop("url")
            self._wml_cpd_creds = _filter_dict(
                cpd_creds.to_dict(),
                [
                    "username",
                    "password",
                    "api_key",
                    "instance_id",
                    "version",
                    "bedrock_url",
                ],
                ["url"],
            )

    def _create_detached_prompt(
        self,
        detached_details: Dict,
        prompt_template_details: Dict,
        detached_asset_details: Dict,
    ) -> str:
        from ibm_aigov_facts_client import (  # type: ignore
            AIGovFactsClient,
            CloudPakforDataConfig,
            DetachedPromptTemplate,
            PromptTemplate,
        )

        try:
            if hasattr(self, "_fact_cpd_creds") and self._fact_cpd_creds:
                cpd_creds = CloudPakforDataConfig(**self._fact_cpd_creds)

                aigov_client = AIGovFactsClient(
                    container_id=self._container_id,
                    container_type=self._container_type,
                    cloud_pak_for_data_configs=cpd_creds,
                    disable_tracing=True,
                )

            else:
                aigov_client = AIGovFactsClient(
                    api_key=self._api_key,
                    container_id=self._container_id,
                    container_type=self._container_type,
                    disable_tracing=True,
                    region=REGIONS_URL[self.region]["factsheet"],
                )

        except Exception as e:
            logging.error(
                f"Error connecting to IBM watsonx.governance (factsheets): {e}",
            )
            raise

        created_detached_pta = aigov_client.assets.create_detached_prompt(
            **detached_asset_details,
            prompt_details=PromptTemplate(**prompt_template_details),
            detached_information=DetachedPromptTemplate(**detached_details),
        )

        return created_detached_pta.to_dict()["asset_id"]

    def _create_deployment_pta(self, asset_id: str, name: str, model_id: str) -> str:
        from ibm_watsonx_ai import APIClient, Credentials  # type: ignore

        try:
            if hasattr(self, "_wml_cpd_creds") and self._wml_cpd_creds:
                creds = Credentials(**self._wml_cpd_creds)

                wml_client = APIClient(creds)
                wml_client.set.default_space(self.space_id)

            else:
                creds = Credentials(
                    url=REGIONS_URL[self.region]["wml"],
                    api_key=self._api_key,
                )
                wml_client = APIClient(creds)
                wml_client.set.default_space(self.space_id)

        except Exception as e:
            logging.error(f"Error connecting to IBM watsonx.ai Runtime: {e}")
            raise

        meta_props = {
            wml_client.deployments.ConfigurationMetaNames.PROMPT_TEMPLATE: {
                "id": asset_id,
            },
            wml_client.deployments.ConfigurationMetaNames.DETACHED: {},
            wml_client.deployments.ConfigurationMetaNames.NAME: name
            + " "
            + "deployment",
            wml_client.deployments.ConfigurationMetaNames.BASE_MODEL_ID: model_id,
        }

        created_deployment = wml_client.deployments.create(asset_id, meta_props)

        return wml_client.deployments.get_uid(created_deployment)

    @deprecated(
        reason="'add_prompt_observer()' is deprecated and will be removed in a future version. Use 'create_prompt_monitor()' instead.",
        version="1.0.5",
        action="always",
    )
    def add_prompt_observer(
        self,
        name: str,
        model_id: str,
        task_id: Literal[
            "extraction",
            "generation",
            "question_answering",
            "retrieval_augmented_generation",
            "summarization",
        ],
        detached_model_provider: str,
        description: str = "",
        model_parameters: Dict = None,
        detached_model_name: str = None,
        detached_model_url: str = None,
        detached_prompt_url: str = None,
        detached_prompt_additional_info: Dict = None,
        prompt_variables: List[str] = None,
        locale: str = "en",
        input_text: str = None,
        context_fields: List[str] = None,
        question_field: str = None,
    ) -> Dict:
        return self.create_prompt_monitor(
            name=name,
            model_id=model_id,
            task_id=task_id,
            detached_model_provider=detached_model_provider,
            description=description,
            model_parameters=model_parameters,
            detached_model_name=detached_model_name,
            detached_model_url=detached_model_url,
            detached_prompt_url=detached_prompt_url,
            detached_prompt_additional_info=detached_prompt_additional_info,
            prompt_variables=prompt_variables,
            locale=locale,
            input_text=input_text,
            context_fields=context_fields,
            question_field=question_field,
        )

    @deprecated(
        reason="'add_prompt_monitor()' is deprecated and will be removed in a future version. Use 'create_prompt_monitor()' instead.",
        version="1.0.6",
        action="always",
    )
    def add_prompt_monitor(
        self,
        name: str,
        model_id: str,
        task_id: Literal[
            "extraction",
            "generation",
            "question_answering",
            "retrieval_augmented_generation",
            "summarization",
        ],
        detached_model_provider: str,
        description: str = "",
        model_parameters: Dict = None,
        detached_model_name: str = None,
        detached_model_url: str = None,
        detached_prompt_url: str = None,
        detached_prompt_additional_info: Dict = None,
        prompt_variables: List[str] = None,
        locale: str = "en",
        input_text: str = None,
        context_fields: List[str] = None,
        question_field: str = None,
    ) -> Dict:
        return self.create_prompt_monitor(
            name=name,
            model_id=model_id,
            task_id=task_id,
            detached_model_provider=detached_model_provider,
            description=description,
            model_parameters=model_parameters,
            detached_model_name=detached_model_name,
            detached_model_url=detached_model_url,
            detached_prompt_url=detached_prompt_url,
            detached_prompt_additional_info=detached_prompt_additional_info,
            prompt_variables=prompt_variables,
            locale=locale,
            input_text=input_text,
            context_fields=context_fields,
            question_field=question_field,
        )

    def create_prompt_monitor(
        self,
        name: str,
        model_id: str,
        task_id: Literal[
            "extraction",
            "generation",
            "question_answering",
            "retrieval_augmented_generation",
            "summarization",
        ],
        detached_model_provider: str,
        description: str = "",
        model_parameters: Dict = None,
        detached_model_name: str = None,
        detached_model_url: str = None,
        detached_prompt_url: str = None,
        detached_prompt_additional_info: Dict = None,
        prompt_variables: List[str] = None,
        locale: str = "en",
        input_text: str = None,
        context_fields: List[str] = None,
        question_field: str = None,
    ) -> Dict:
        """
        Creates a detached (external) prompt template asset and attaches a monitor to the specified prompt template asset.

        Args:
            name (str): The name of the External Prompt Template Asset.
            model_id (str): The ID of the model associated with the prompt.
            task_id (str): The task identifier.
            detached_model_provider (str): The external model provider.
            description (str, optional): A description of the External Prompt Template Asset.
            model_parameters (Dict, optional): Model parameters and their respective values.
            detached_model_name (str, optional): The name of the external model.
            detached_model_url (str, optional): The URL of the external model.
            detached_prompt_url (str, optional): The URL of the external prompt.
            detached_prompt_additional_info (Dict, optional): Additional information related to the external prompt.
            prompt_variables (List[str], optional): Values for the prompt variables.
            locale (str, optional): Locale code for the input/output language. eg. "en", "pt", "es".
            input_text (str, optional): The input text for the prompt.
            context_fields (List[str], optional): A list of fields that will provide context to the prompt.
                Applicable only for "retrieval_augmented_generation" task type.
            question_field (str, optional): The field containing the question to be answered.
                Applicable only for "retrieval_augmented_generation" task type.

        Example:
            ```python
            wxgov_client.create_prompt_monitor(
                name="Detached prompt (model AWS Anthropic)",
                model_id="anthropic.claude-v2",
                task_id="retrieval_augmented_generation",
                detached_model_provider="AWS Bedrock",
                detached_model_name="Anthropic Claude 2.0",
                detached_model_url="https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-claude.html",
                prompt_variables=["context1", "context2", "input_query"],
                input_text="Prompt text to be given",
                context_fields=["context1", "context2"],
                question_field="input_query",
            )
            ```
        """
        if (not (self.project_id or self.space_id)) or (
            self.project_id and self.space_id
        ):
            raise ValueError(
                "Invalid configuration: Neither was provided: please set either 'project_id' or 'space_id'. "
                "Both were provided: 'project_id' and 'space_id' cannot be set at the same time."
            )

        if task_id == "retrieval_augmented_generation":
            if not context_fields or not question_field:
                raise ValueError(
                    "For 'retrieval_augmented_generation' task, requires non-empty 'context_fields' and 'question_field'."
                )

        prompt_metadata = locals()
        # Remove unused vars from dict
        prompt_metadata.pop("self", None)
        prompt_metadata.pop("context_fields", None)
        prompt_metadata.pop("question_field", None)
        prompt_metadata.pop("locale", None)

        # Update name of keys to aigov_facts api
        prompt_metadata["input"] = prompt_metadata.pop("input_text", None)
        prompt_metadata["model_provider"] = prompt_metadata.pop(
            "detached_model_provider",
            None,
        )
        prompt_metadata["model_name"] = prompt_metadata.pop("detached_model_name", None)
        prompt_metadata["model_url"] = prompt_metadata.pop("detached_model_url", None)
        prompt_metadata["prompt_url"] = prompt_metadata.pop("detached_prompt_url", None)
        prompt_metadata["prompt_additional_info"] = prompt_metadata.pop(
            "detached_prompt_additional_info",
            None,
        )

        # Update list of vars to dict
        prompt_metadata["prompt_variables"] = Dict.fromkeys(
            prompt_metadata["prompt_variables"], ""
        )

        from ibm_watson_openscale import APIClient as WosAPIClient  # type: ignore

        if not self._wos_client:
            try:
                if hasattr(self, "_wos_cpd_creds") and self._wos_cpd_creds:
                    from ibm_cloud_sdk_core.authenticators import (
                        CloudPakForDataAuthenticator,  # type: ignore
                    )

                    authenticator = CloudPakForDataAuthenticator(**self._wos_cpd_creds)
                    self._wos_client = WosAPIClient(
                        authenticator=authenticator,
                        service_url=self._wos_cpd_creds["url"],
                    )

                else:
                    from ibm_cloud_sdk_core.authenticators import (
                        IAMAuthenticator,  # type: ignore
                    )

                    authenticator = IAMAuthenticator(apikey=self._api_key)
                    self._wos_client = WosAPIClient(
                        authenticator=authenticator,
                        service_url=REGIONS_URL[self.region]["wos"],
                    )

            except Exception as e:
                logging.error(
                    f"Error connecting to IBM watsonx.governance (openscale): {e}",
                )
                raise

        detached_details = _filter_dict(
            prompt_metadata,
            ["model_name", "model_url", "prompt_url", "prompt_additional_info"],
            ["model_id", "model_provider"],
        )
        detached_details["prompt_id"] = "detached_prompt_" + str(uuid.uuid4())

        prompt_details = _filter_dict(
            prompt_metadata,
            ["prompt_variables", "input", "model_parameters"],
        )

        detached_asset_details = _filter_dict(
            prompt_metadata,
            ["description"],
            ["name", "model_id", "task_id"],
        )

        detached_pta_id = suppress_output(
            self._create_detached_prompt,
            detached_details,
            prompt_details,
            detached_asset_details,
        )
        deployment_id = None
        if self._container_type == "space":
            deployment_id = suppress_output(
                self._create_deployment_pta, detached_pta_id, name, model_id
            )

        monitors = {
            "generative_ai_quality": {
                "parameters": {"min_sample_size": 10, "metrics_configuration": {}},
            },
        }

        max_attempt_execute_prompt_setup = 0
        while max_attempt_execute_prompt_setup < 2:
            try:
                generative_ai_monitor_details = suppress_output(
                    self._wos_client.wos.execute_prompt_setup,
                    prompt_template_asset_id=detached_pta_id,
                    space_id=self.space_id,
                    project_id=self.project_id,
                    deployment_id=deployment_id,
                    label_column="reference_output",
                    context_fields=context_fields,
                    question_field=question_field,
                    operational_space_id=self._deployment_stage,
                    problem_type=task_id,
                    data_input_locale=[locale],
                    generated_output_locale=[locale],
                    input_data_type="unstructured_text",
                    supporting_monitors=monitors,
                    background_mode=False,
                )

                break

            except Exception as e:
                if (
                    e.code == 403
                    and "The user entitlement does not exist" in e.message
                    and max_attempt_execute_prompt_setup < 1
                ):
                    max_attempt_execute_prompt_setup = (
                        max_attempt_execute_prompt_setup + 1
                    )

                    data_marts = self._wos_client.data_marts.list().result
                    if (data_marts.data_marts is None) or (not data_marts.data_marts):
                        raise ValueError(
                            "Error retrieving IBM watsonx.governance (openscale) data mart. \
                                         Make sure the data mart are configured.",
                        )

                    data_mart_id = data_marts.data_marts[0].metadata.id

                    self._wos_client.wos.add_instance_mapping(
                        service_instance_id=data_mart_id,
                        space_id=self.space_id,
                        project_id=self.project_id,
                    )
                else:
                    max_attempt_execute_prompt_setup = 2
                    raise

        generative_ai_monitor_details = generative_ai_monitor_details.result._to_dict()

        return {
            "detached_prompt_template_asset_id": detached_pta_id,
            "deployment_id": deployment_id,
            "subscription_id": generative_ai_monitor_details["subscription_id"],
        }

    def store_payload_records(
        self,
        request_records: List[Dict],
        subscription_id: str = None,
    ) -> List[str]:
        """
        Stores records to the payload logging system.

        Args:
            request_records (List[Dict]): A list of records to be logged, where each record is represented as a dictionary.
            subscription_id (str, optional): The subscription ID associated with the records being logged.

        Example:
            ```python
            wxgov_client.store_payload_records(
                request_records=[
                    {
                        "context1": "value_context1",
                        "context2": "value_context1",
                        "input_query": "What's Beekeeper Framework?",
                        "generated_text": "Beekeeper is a data framework to make AI easier to work with.",
                        "input_token_count": 25,
                        "generated_token_count": 150,
                    }
                ],
                subscription_id="5d62977c-a53d-4b6d-bda1-7b79b3b9d1a0",
            )
            ```
        """
        from ibm_cloud_sdk_core.authenticators import IAMAuthenticator
        from ibm_watson_openscale import APIClient as WosAPIClient
        from ibm_watson_openscale.supporting_classes.enums import (
            DataSetTypes,
            TargetTypes,
        )

        # Expected behavior: Prefer using fn `subscription_id`.
        # Fallback to `self.subscription_id` if `subscription_id` None or empty.
        _subscription_id = subscription_id or self.subscription_id

        if _subscription_id is None or _subscription_id == "":
            raise ValueError(
                "Unexpected value for 'subscription_id': Cannot be None or empty string."
            )

        if not self._wos_client:
            try:
                if hasattr(self, "_wos_cpd_creds") and self._wos_cpd_creds:
                    from ibm_cloud_sdk_core.authenticators import (
                        CloudPakForDataAuthenticator,  # type: ignore
                    )

                    authenticator = CloudPakForDataAuthenticator(**self._wos_cpd_creds)
                    self._wos_client = WosAPIClient(
                        authenticator=authenticator,
                        service_url=self._wos_cpd_creds["url"],
                    )

                else:
                    from ibm_cloud_sdk_core.authenticators import (
                        IAMAuthenticator,  # type: ignore
                    )

                    authenticator = IAMAuthenticator(apikey=self._api_key)
                    self._wos_client = WosAPIClient(
                        authenticator=authenticator,
                        service_url=REGIONS_URL[self.region]["wos"],
                    )

            except Exception as e:
                logging.error(
                    f"Error connecting to IBM watsonx.governance (openscale): {e}",
                )
                raise

        subscription_details = self._wos_client.subscriptions.get(
            _subscription_id,
        ).result
        subscription_details = json.loads(str(subscription_details))

        feature_fields = subscription_details["entity"]["asset_properties"][
            "feature_fields"
        ]

        payload_data_set_id = (
            self._wos_client.data_sets.list(
                type=DataSetTypes.PAYLOAD_LOGGING,
                target_target_id=_subscription_id,
                target_target_type=TargetTypes.SUBSCRIPTION,
            )
            .result.data_sets[0]
            .metadata.id
        )

        payload_data = _convert_payload_format(request_records, feature_fields)

        suppress_output(
            self._wos_client.data_sets.store_records,
            data_set_id=payload_data_set_id,
            request_body=payload_data,
            background_mode=False,
        )

        return [data["scoring_id"] + "-1" for data in payload_data]

    def __call__(self, payload: PayloadRecord) -> None:
        if self.prompt_template:
            template_vars = extract_template_vars(
                self.prompt_template.template, payload.input_text
            )

        if not template_vars:
            self.store_payload_records([payload.model_dump()])
        else:
            self.store_payload_records([{**payload.model_dump(), **template_vars}])


class WatsonxPromptMonitor(PromptMonitor):
    """
    Provides functionality to interact with IBM watsonx.governance for monitoring prompts executed within
    IBM watsonx.ai LLMs.

    Note:
        One of the following parameters is required to create a prompt monitor:
        `project_id` or `space_id`, but not both.

    Attributes:
        api_key (str): The API key for IBM watsonx.governance.
        space_id (str, optional): The space ID in watsonx.governance.
        project_id (str, optional): The project ID in watsonx.governance.
        region (str, optional): The region where watsonx.governance is hosted when using IBM Cloud.
            Defaults to `us-south`.
        cpd_creds (CloudPakforDataCredentials, optional): The Cloud Pak for Data environment credentials.
        subscription_id (str, optional): The subscription ID associated with the records being logged.

    Example:
        ```python
        from beekeeper.monitors.watsonx import (
            WatsonxPromptMonitor,
            CloudPakforDataCredentials,
        )

        # watsonx.governance (IBM Cloud)
        wxgov_client = WatsonxPromptMonitor(api_key="API_KEY", space_id="SPACE_ID")

        # watsonx.governance (CP4D)
        cpd_creds = CloudPakforDataCredentials(
            url="CPD_URL",
            username="USERNAME",
            password="PASSWORD",
            version="5.0",
            instance_id="openshift",
        )

        wxgov_client = WatsonxPromptMonitor(space_id="SPACE_ID", cpd_creds=cpd_creds)
        ```
    """

    def __init__(
        self,
        api_key: str = None,
        space_id: str = None,
        project_id: str = None,
        region: Literal["us-south", "eu-de", "au-syd"] = "us-south",
        cpd_creds: CloudPakforDataCredentials | Dict = None,
        subscription_id: str = None,
        **kwargs,
    ) -> None:
        import ibm_aigov_facts_client  # noqa: F401
        import ibm_cloud_sdk_core.authenticators  # noqa: F401
        import ibm_watson_openscale  # noqa: F401
        import ibm_watsonx_ai  # noqa: F401

        super().__init__(**kwargs)

        self.space_id = space_id
        self.project_id = project_id
        self.region = region
        self.subscription_id = subscription_id
        self._api_key = api_key
        self._wos_client = None

        self._container_id = space_id if space_id else project_id
        self._container_type = "space" if space_id else "project"
        self._deployment_stage = "production" if space_id else "development"

        if cpd_creds:
            self._wos_cpd_creds = _filter_dict(
                cpd_creds.to_dict(),
                ["username", "password", "api_key", "disable_ssl_verification"],
                ["url"],
            )
            self._fact_cpd_creds = _filter_dict(
                cpd_creds.to_dict(),
                ["username", "password", "api_key", "bedrock_url"],
                ["url"],
            )
            self._fact_cpd_creds["service_url"] = self._fact_cpd_creds.pop("url")
            self._wml_cpd_creds = _filter_dict(
                cpd_creds.to_dict(),
                [
                    "username",
                    "password",
                    "api_key",
                    "instance_id",
                    "version",
                    "bedrock_url",
                ],
                ["url"],
            )

    def _create_prompt_template(
        self,
        prompt_template_details: Dict,
        asset_details: Dict,
    ) -> str:
        from ibm_aigov_facts_client import (
            AIGovFactsClient,
            CloudPakforDataConfig,
            PromptTemplate,
        )

        try:
            if hasattr(self, "_fact_cpd_creds") and self._fact_cpd_creds:
                cpd_creds = CloudPakforDataConfig(**self._fact_cpd_creds)

                aigov_client = AIGovFactsClient(
                    container_id=self._container_id,
                    container_type=self._container_type,
                    cloud_pak_for_data_configs=cpd_creds,
                    disable_tracing=True,
                )

            else:
                aigov_client = AIGovFactsClient(
                    api_key=self._api_key,
                    container_id=self._container_id,
                    container_type=self._container_type,
                    disable_tracing=True,
                    region=REGIONS_URL[self.region]["factsheet"],
                )

        except Exception as e:
            logging.error(
                f"Error connecting to IBM watsonx.governance (factsheets): {e}",
            )
            raise

        created_pta = aigov_client.assets.create_prompt(
            **asset_details,
            input_mode="freeform",
            prompt_details=PromptTemplate(**prompt_template_details),
        )

        return created_pta.to_dict()["asset_id"]

    def _create_deployment_pta(self, asset_id: str, name: str, model_id: str) -> str:
        from ibm_watsonx_ai import APIClient, Credentials  # type: ignore

        try:
            if hasattr(self, "_wml_cpd_creds") and self._wml_cpd_creds:
                creds = Credentials(**self._wml_cpd_creds)

                wml_client = APIClient(creds)
                wml_client.set.default_space(self.space_id)

            else:
                creds = Credentials(
                    url=REGIONS_URL[self.region]["wml"],
                    api_key=self._api_key,
                )

                wml_client = APIClient(creds)
                wml_client.set.default_space(self.space_id)

        except Exception as e:
            logging.error(f"Error connecting to IBM watsonx.ai Runtime: {e}")
            raise

        meta_props = {
            wml_client.deployments.ConfigurationMetaNames.PROMPT_TEMPLATE: {
                "id": asset_id,
            },
            wml_client.deployments.ConfigurationMetaNames.FOUNDATION_MODEL: {},
            wml_client.deployments.ConfigurationMetaNames.NAME: name
            + " "
            + "deployment",
            wml_client.deployments.ConfigurationMetaNames.BASE_MODEL_ID: model_id,
        }

        created_deployment = wml_client.deployments.create(asset_id, meta_props)

        return wml_client.deployments.get_uid(created_deployment)

    @deprecated(
        reason="'add_prompt_observer()' is deprecated and will be removed in a future version. Use 'create_prompt_monitor()' instead.",
        version="1.0.5",
        action="always",
    )
    def add_prompt_observer(
        self,
        name: str,
        model_id: str,
        task_id: Literal[
            "extraction",
            "generation",
            "question_answering",
            "retrieval_augmented_generation",
            "summarization",
        ],
        description: str = "",
        model_parameters: Dict = None,
        prompt_variables: List[str] = None,
        locale: str = "en",
        input_text: str = None,
        context_fields: List[str] = None,
        question_field: str = None,
    ) -> Dict:
        return self.create_prompt_monitor(
            name=name,
            model_id=model_id,
            task_id=task_id,
            description=description,
            model_parameters=model_parameters,
            prompt_variables=prompt_variables,
            locale=locale,
            input_text=input_text,
            context_fields=context_fields,
            question_field=question_field,
        )

    @deprecated(
        reason="'add_prompt_observer()' is deprecated and will be removed in a future version. Use 'create_prompt_monitor()' instead.",
        version="1.0.6",
        action="always",
    )
    def add_prompt_monitor(
        self,
        name: str,
        model_id: str,
        task_id: Literal[
            "extraction",
            "generation",
            "question_answering",
            "retrieval_augmented_generation",
            "summarization",
        ],
        description: str = "",
        model_parameters: Dict = None,
        prompt_variables: List[str] = None,
        locale: str = "en",
        input_text: str = None,
        context_fields: List[str] = None,
        question_field: str = None,
    ) -> Dict:
        return self.create_prompt_monitor(
            name=name,
            model_id=model_id,
            task_id=task_id,
            description=description,
            model_parameters=model_parameters,
            prompt_variables=prompt_variables,
            locale=locale,
            input_text=input_text,
            context_fields=context_fields,
            question_field=question_field,
        )

    def create_prompt_monitor(
        self,
        name: str,
        model_id: str,
        task_id: Literal[
            "extraction",
            "generation",
            "question_answering",
            "retrieval_augmented_generation",
            "summarization",
        ],
        description: str = "",
        model_parameters: Dict = None,
        prompt_variables: List[str] = None,
        locale: str = "en",
        input_text: str = None,
        context_fields: List[str] = None,
        question_field: str = None,
    ) -> Dict:
        """
        Creates an IBM Prompt Template Asset and ssetup monitor for the given prompt template asset.

        Args:
            name (str): The name of the Prompt Template Asset.
            model_id (str): The ID of the model associated with the prompt.
            task_id (str): The task identifier.
            description (str, optional): A description of the Prompt Template Asset.
            model_parameters (Dict, optional): A dictionary of model parameters and their respective values.
            prompt_variables (List[str], optional): A list of values for prompt input variables.
            locale (str, optional): Locale code for the input/output language. eg. "en", "pt", "es".
            input_text (str, optional): The input text for the prompt.
            context_fields (List[str], optional): A list of fields that will provide context to the prompt.
                Applicable only for the `retrieval_augmented_generation` task type.
            question_field (str, optional): The field containing the question to be answered.
                Applicable only for the `retrieval_augmented_generation` task type.

        Example:
            ```python
            wxgov_client.create_prompt_monitor(
                name="IBM prompt template",
                model_id="ibm/granite-3-2b-instruct",
                task_id="retrieval_augmented_generation",
                prompt_variables=["context1", "context2", "input_query"],
                input_text="Prompt text to be given",
                context_fields=["context1", "context2"],
                question_field="input_query",
            )
            ```
        """
        if (not (self.project_id or self.space_id)) or (
            self.project_id and self.space_id
        ):
            raise ValueError(
                "Invalid configuration: Neither was provided: please set either 'project_id' or 'space_id'. "
                "Both were provided: 'project_id' and 'space_id' cannot be set at the same time."
            )

        if task_id == "retrieval_augmented_generation":
            if not context_fields or not question_field:
                raise ValueError(
                    "For 'retrieval_augmented_generation' task, requires non-empty 'context_fields' and 'question_field'."
                )

        prompt_metadata = locals()
        # Remove unused vars from dict
        prompt_metadata.pop("self", None)
        prompt_metadata.pop("context_fields", None)
        prompt_metadata.pop("question_field", None)
        prompt_metadata.pop("locale", None)

        # Update name of keys to aigov_facts api
        prompt_metadata["input"] = prompt_metadata.pop("input_text", None)

        # Update list of vars to dict
        prompt_metadata["prompt_variables"] = Dict.fromkeys(
            prompt_metadata["prompt_variables"], ""
        )

        from ibm_cloud_sdk_core.authenticators import IAMAuthenticator  # type: ignore
        from ibm_watson_openscale import APIClient as WosAPIClient  # type: ignore

        if not self._wos_client:
            try:
                if hasattr(self, "_wos_cpd_creds") and self._wos_cpd_creds:
                    from ibm_cloud_sdk_core.authenticators import (
                        CloudPakForDataAuthenticator,  # type: ignore
                    )

                    authenticator = CloudPakForDataAuthenticator(**self._wos_cpd_creds)

                    self._wos_client = WosAPIClient(
                        authenticator=authenticator,
                        service_url=self._wos_cpd_creds["url"],
                    )

                else:
                    from ibm_cloud_sdk_core.authenticators import (
                        IAMAuthenticator,  # type: ignore
                    )

                    authenticator = IAMAuthenticator(apikey=self._api_key)
                    self._wos_client = WosAPIClient(
                        authenticator=authenticator,
                        service_url=REGIONS_URL[self.region]["wos"],
                    )

            except Exception as e:
                logging.error(
                    f"Error connecting to IBM watsonx.governance (openscale): {e}",
                )
                raise

        prompt_details = _filter_dict(
            prompt_metadata,
            ["prompt_variables", "input", "model_parameters"],
        )

        asset_details = _filter_dict(
            prompt_metadata,
            ["description"],
            ["name", "model_id", "task_id"],
        )

        pta_id = suppress_output(
            self._create_prompt_template, prompt_details, asset_details
        )
        deployment_id = None
        if self._container_type == "space":
            deployment_id = suppress_output(
                self._create_deployment_pta, pta_id, name, model_id
            )

        monitors = {
            "generative_ai_quality": {
                "parameters": {"min_sample_size": 10, "metrics_configuration": {}},
            },
        }

        max_attempt_execute_prompt_setup = 0
        while max_attempt_execute_prompt_setup < 2:
            try:
                generative_ai_monitor_details = suppress_output(
                    self._wos_client.wos.execute_prompt_setup,
                    prompt_template_asset_id=pta_id,
                    space_id=self.space_id,
                    project_id=self.project_id,
                    deployment_id=deployment_id,
                    label_column="reference_output",
                    context_fields=context_fields,
                    question_field=question_field,
                    operational_space_id=self._deployment_stage,
                    problem_type=task_id,
                    data_input_locale=[locale],
                    generated_output_locale=[locale],
                    input_data_type="unstructured_text",
                    supporting_monitors=monitors,
                    background_mode=False,
                ).result

                break

            except Exception as e:
                if (
                    e.code == 403
                    and "The user entitlement does not exist" in e.message
                    and max_attempt_execute_prompt_setup < 1
                ):
                    max_attempt_execute_prompt_setup = (
                        max_attempt_execute_prompt_setup + 1
                    )

                    data_marts = self._wos_client.data_marts.list().result
                    if (data_marts.data_marts is None) or (not data_marts.data_marts):
                        raise ValueError(
                            "Error retrieving IBM watsonx.governance (openscale) data mart. \
                                         Make sure the data mart are configured.",
                        )

                    data_mart_id = data_marts.data_marts[0].metadata.id

                    self._wos_client.wos.add_instance_mapping(
                        service_instance_id=data_mart_id,
                        space_id=self.space_id,
                        project_id=self.project_id,
                    )
                else:
                    max_attempt_execute_prompt_setup = 2
                    raise

        generative_ai_monitor_details = generative_ai_monitor_details._to_dict()

        return {
            "prompt_template_asset_id": pta_id,
            "deployment_id": deployment_id,
            "subscription_id": generative_ai_monitor_details["subscription_id"],
        }

    def store_payload_records(
        self,
        request_records: List[Dict],
        subscription_id: str = None,
    ) -> List[str]:
        """
        Stores records to the payload logging system.

        Args:
            request_records (List[Dict]): A list of records to be logged. Each record is represented as a dictionary.
            subscription_id (str, optional): The subscription ID associated with the records being logged.

        Example:
            ```python
            wxgov_client.store_payload_records(
                request_records=[
                    {
                        "context1": "value_context1",
                        "context2": "value_context1",
                        "input_query": "What's Beekeeper Framework?",
                        "generated_text": "Beekeeper is a data framework to make AI easier to work with.",
                        "input_token_count": 25,
                        "generated_token_count": 150,
                    }
                ],
                subscription_id="5d62977c-a53d-4b6d-bda1-7b79b3b9d1a0",
            )
            ```
        """
        from ibm_cloud_sdk_core.authenticators import IAMAuthenticator
        from ibm_watson_openscale import APIClient as WosAPIClient
        from ibm_watson_openscale.supporting_classes.enums import (
            DataSetTypes,
            TargetTypes,
        )

        # Expected behavior: Prefer using fn `subscription_id`.
        # Fallback to `self.subscription_id` if `subscription_id` None or empty.
        _subscription_id = subscription_id or self.subscription_id

        if _subscription_id is None or _subscription_id == "":
            raise ValueError(
                "Unexpected value for 'subscription_id': Cannot be None or empty string."
            )

        if not self._wos_client:
            try:
                if hasattr(self, "_wos_cpd_creds") and self._wos_cpd_creds:
                    from ibm_cloud_sdk_core.authenticators import (
                        CloudPakForDataAuthenticator,  # type: ignore
                    )

                    authenticator = CloudPakForDataAuthenticator(**self._wos_cpd_creds)

                    self._wos_client = WosAPIClient(
                        authenticator=authenticator,
                        service_url=self._wos_cpd_creds["url"],
                    )

                else:
                    from ibm_cloud_sdk_core.authenticators import (
                        IAMAuthenticator,  # type: ignore
                    )

                    authenticator = IAMAuthenticator(apikey=self._api_key)
                    self._wos_client = WosAPIClient(
                        authenticator=authenticator,
                        service_url=REGIONS_URL[self.region]["wos"],
                    )

            except Exception as e:
                logging.error(
                    f"Error connecting to IBM watsonx.governance (openscale): {e}",
                )
                raise

        subscription_details = self._wos_client.subscriptions.get(
            _subscription_id,
        ).result
        subscription_details = json.loads(str(subscription_details))

        feature_fields = subscription_details["entity"]["asset_properties"][
            "feature_fields"
        ]

        payload_data_set_id = (
            self._wos_client.data_sets.list(
                type=DataSetTypes.PAYLOAD_LOGGING,
                target_target_id=_subscription_id,
                target_target_type=TargetTypes.SUBSCRIPTION,
            )
            .result.data_sets[0]
            .metadata.id
        )

        payload_data = _convert_payload_format(request_records, feature_fields)

        suppress_output(
            self._wos_client.data_sets.store_records,
            data_set_id=payload_data_set_id,
            request_body=payload_data,
            background_mode=False,
        )

        return [data["scoring_id"] + "-1" for data in payload_data]

    def __call__(self, payload: PayloadRecord) -> None:
        if self.prompt_template:
            template_vars = extract_template_vars(
                self.prompt_template.template, payload.input_text
            )

        if not template_vars:
            self.store_payload_records([payload.model_dump()])
        else:
            self.store_payload_records([{**payload.model_dump(), **template_vars}])


# ===== Supporting Classes =====
class WatsonxLocalMetric(BaseModel):
    """
    Provides the IBM watsonx.governance local monitor metric definition.

    Attributes:
        name (str): The name of the metric.
        data_type (str): The data type of the metric. Currently supports "string", "integer", "double", and "timestamp".
        nullable (bool, optional): Indicates whether the metric can be null. Defaults to `False`.

    Example:
        ```python
        from beekeeper.monitors.watsonx import WatsonxLocalMetric

        WatsonxLocalMetric(name="context_quality", data_type="double")
        ```
    """

    name: str
    data_type: Literal["string", "integer", "double", "timestamp"]
    nullable: bool = True

    def to_dict(self) -> Dict:
        return {"name": self.name, "type": self.data_type, "nullable": self.nullable}


class WatsonxMetricThreshold(BaseModel):
    """
    Defines the metric threshold for IBM watsonx.governance.

    Attributes:
        threshold_type (str): The threshold type. Can be either `lower_limit` or `upper_limit`.
        default_value (float): The metric threshold value.

    Example:
        ```python
        from beekeeper.monitors.watsonx import WatsonxMetricThreshold

        WatsonxMetricThreshold(threshold_type="lower_limit", default_value=0.8)
        ```
    """

    threshold_type: Literal["lower_limit", "upper_limit"]
    default_value: float = None

    def to_dict(self) -> Dict:
        return {"type": self.threshold_type, "default": self.default_value}


class WatsonxMetric(BaseModel):
    """
    Defines the IBM watsonx.governance global monitor metric.

    Attributes:
        name (str): The name of the metric.
        applies_to (List[str]): A list of task types that the metric applies to. Currently supports:
            "summarization", "generation", "question_answering", "extraction", and "retrieval_augmented_generation".
        thresholds (List[WatsonxMetricThreshold]): A list of metric thresholds associated with the metric.

    Example:
        ```python
        from beekeeper.monitors.watsonx import (
            WatsonxMetric,
            WatsonxMetricThreshold,
        )

        WatsonxMetric(
            name="context_quality",
            applies_to=["retrieval_augmented_generation", "summarization"],
            thresholds=[
                WatsonxMetricThreshold(threshold_type="lower_limit", default_value=0.75)
            ],
        )
        ```
    """

    name: str
    applies_to: List[
        Literal[
            "summarization",
            "generation",
            "question_answering",
            "extraction",
            "retrieval_augmented_generation",
        ]
    ]
    thresholds: Optional[List[WatsonxMetricThreshold]] = None

    def to_dict(self) -> Dict:
        from ibm_watson_openscale.base_classes.watson_open_scale_v2 import (
            ApplicabilitySelection,
            MetricThreshold,
        )

        monitor_metric = {
            "name": self.name,
            "applies_to": ApplicabilitySelection(problem_type=self.applies_to),
        }

        if self.thresholds is not None:
            monitor_metric["thresholds"] = [
                MetricThreshold(**threshold.to_dict()) for threshold in self.thresholds
            ]

        return monitor_metric


# ===== Metric Classes =====


class WatsonxCustomMetricsManager:
    """
    Provides functionality to set up a custom metric to measure your model's performance with IBM watsonx.governance.

    Attributes:
        api_key (str): The API key for IBM watsonx.governance.
        region (str, optional): The region where IBM watsonx.governance is hosted when using IBM Cloud.
            Defaults to `us-south`.
        cpd_creds (CloudPakforDataCredentials, optional): IBM Cloud Pak for Data environment credentials.

    Example:
        ```python
        from beekeeper.monitors.watsonx import (
            WatsonxCustomMetricsManager,
            CloudPakforDataCredentials,
        )

        # watsonx.governance (IBM Cloud)
        wxgov_client = WatsonxCustomMetricsManager(api_key="API_KEY")

        # watsonx.governance (CP4D)
        cpd_creds = CloudPakforDataCredentials(
            url="CPD_URL",
            username="USERNAME",
            password="PASSWORD",
            version="5.0",
            instance_id="openshift",
        )

        wxgov_client = WatsonxCustomMetricsManager(cpd_creds=cpd_creds)
        ```
    """

    def __init__(
        self,
        api_key: str = None,
        region: Literal["us-south", "eu-de", "au-syd"] = "us-south",
        cpd_creds: CloudPakforDataCredentials | Dict = None,
    ) -> None:
        from ibm_cloud_sdk_core.authenticators import IAMAuthenticator  # type: ignore
        from ibm_watson_openscale import APIClient as WosAPIClient  # type: ignore

        self.region = region
        self._api_key = api_key
        self._wos_client = None

        if cpd_creds:
            self._wos_cpd_creds = _filter_dict(
                cpd_creds.to_dict(),
                ["username", "password", "api_key", "disable_ssl_verification"],
                ["url"],
            )

        if not self._wos_client:
            try:
                if hasattr(self, "_wos_cpd_creds") and self._wos_cpd_creds:
                    from ibm_cloud_sdk_core.authenticators import (
                        CloudPakForDataAuthenticator,  # type: ignore
                    )

                    authenticator = CloudPakForDataAuthenticator(**self._wos_cpd_creds)

                    self._wos_client = WosAPIClient(
                        authenticator=authenticator,
                        service_url=self._wos_cpd_creds["url"],
                    )

                else:
                    from ibm_cloud_sdk_core.authenticators import (
                        IAMAuthenticator,  # type: ignore
                    )

                    authenticator = IAMAuthenticator(apikey=self._api_key)
                    self._wos_client = WosAPIClient(
                        authenticator=authenticator,
                        service_url=REGIONS_URL[self.region]["wos"],
                    )

            except Exception as e:
                logging.error(
                    f"Error connecting to IBM watsonx.governance (openscale): {e}",
                )
                raise

    def _add_integrated_system(
        self,
        credentials: IntegratedSystemCredentials,
        name: str,
        endpoint: str,
    ) -> str:
        custom_metrics_integrated_system = self._wos_client.integrated_systems.add(
            name=name,
            description="Integrated system created by Beekeeper.",
            type="custom_metrics_provider",
            credentials=credentials.to_dict(),
            connection={"display_name": name, "endpoint": endpoint},
        ).result

        return custom_metrics_integrated_system.metadata.id

    def _add_monitor_definitions(
        self,
        name: str,
        metrics: List[WatsonxMetric],
        schedule: bool,
    ):
        from ibm_watson_openscale.base_classes.watson_open_scale_v2 import (
            ApplicabilitySelection,
            MonitorInstanceSchedule,
            MonitorMetricRequest,
            MonitorRuntime,
            ScheduleStartTime,
        )

        _metrics = [MonitorMetricRequest(**metric.to_dict()) for metric in metrics]
        _monitor_runtime = None
        _monitor_schedule = None

        if schedule:
            _monitor_runtime = MonitorRuntime(type="custom_metrics_provider")
            _monitor_schedule = MonitorInstanceSchedule(
                repeat_interval=1,
                repeat_unit="hour",
                start_time=ScheduleStartTime(
                    type="relative",
                    delay_unit="minute",
                    delay=30,
                ),
            )

        custom_monitor_details = self._wos_client.monitor_definitions.add(
            name=name,
            metrics=_metrics,
            tags=[],
            schedule=_monitor_schedule,
            applies_to=ApplicabilitySelection(input_data_type=["unstructured_text"]),
            monitor_runtime=_monitor_runtime,
            background_mode=False,
        ).result

        return custom_monitor_details.metadata.id

    def _get_monitor_instance(self, subscription_id: str, monitor_definition_id: str):
        monitor_instances = self._wos_client.monitor_instances.list(
            monitor_definition_id=monitor_definition_id,
            target_target_id=subscription_id,
        ).result.monitor_instances

        if len(monitor_instances) == 1:
            return monitor_instances[0]
        else:
            return None

    def _update_monitor_instance(
        self,
        integrated_system_id: str,
        custom_monitor_id: str,
    ):
        payload = [
            {
                "op": "replace",
                "path": "/parameters",
                "value": {
                    "custom_metrics_provider_id": integrated_system_id,
                    "custom_metrics_wait_time": 60,
                    "enable_custom_metric_runs": True,
                },
            },
        ]

        return self._wos_client.monitor_instances.update(
            custom_monitor_id,
            payload,
            update_metadata_only=True,
        ).result

    def _get_patch_request_field(
        self,
        field_path: str,
        field_value: Any,
        op_name: str = "replace",
    ) -> Dict:
        return {"op": op_name, "path": field_path, "value": field_value}

    def _get_dataset_id(
        self,
        subscription_id: str,
        data_set_type: Literal["feedback", "payload_logging"],
    ) -> str:
        data_sets = self._wos_client.data_sets.list(
            target_target_id=subscription_id,
            type=data_set_type,
        ).result.data_sets
        data_set_id = None
        if len(data_sets) > 0:
            data_set_id = data_sets[0].metadata.id
        return data_set_id

    def _get_dataset_data(self, data_set_id: str):
        json_data = self._wos_client.data_sets.get_list_of_records(
            data_set_id=data_set_id,
            format="list",
        ).result

        if not json_data.get("records"):
            return None

        return json_data["records"][0]

    def _get_existing_data_mart(self):
        data_marts = self._wos_client.data_marts.list().result.data_marts
        if len(data_marts) == 0:
            raise Exception(
                "No data marts found. Please ensure at least one data mart is available.",
            )

        return data_marts[0].metadata.id

    # ===== Global Custom Metrics =====
    @deprecated(
        reason="'add_metric_definition()' is deprecated and will be removed in a future version. Use 'create_metric_definition()' instead.",
        version="1.0.6",
        action="always",
    )
    def add_metric_definition(
        self,
        name: str,
        metrics: List[WatsonxMetric],
        integrated_system_url: str,
        integrated_system_credentials: IntegratedSystemCredentials,
        schedule: bool = False,
    ):
        return self.create_metric_definition(
            name=name,
            metrics=metrics,
            integrated_system_url=integrated_system_url,
            integrated_system_credentials=integrated_system_credentials,
            schedule=schedule,
        )

    def create_metric_definition(
        self,
        name: str,
        metrics: List[WatsonxMetric],
        integrated_system_url: str,
        integrated_system_credentials: IntegratedSystemCredentials,
        schedule: bool = False,
    ):
        """
        Creates a custom metric definition for IBM watsonx.governance.

        This must be done before using custom metrics.

        Args:
            name (str): The name of the custom metric group.
            metrics (List[WatsonxMetric]): A list of metrics to be measured.
            schedule (bool, optional): Enable or disable the scheduler. Defaults to `False`.
            integrated_system_url (str): The URL of the external metric provider.
            integrated_system_credentials (IntegratedSystemCredentials): The credentials for the integrated system.

        Example:
            ```python
            from beekeeper.monitors.watsonx import (
                WatsonxMetric,
                IntegratedSystemCredentials,
                WatsonxMetricThreshold,
            )

            wxgov_client.create_metric_definition(
                name="Custom Metric - Custom LLM Quality",
                metrics=[
                    WatsonxMetric(
                        name="context_quality",
                        applies_to=[
                            "retrieval_augmented_generation",
                            "summarization",
                        ],
                        thresholds=[
                            WatsonxMetricThreshold(
                                threshold_type="lower_limit", default_value=0.75
                            )
                        ],
                    )
                ],
                integrated_system_url="IS_URL",  # URL to the endpoint computing the metric
                integrated_system_credentials=IntegratedSystemCredentials(
                    auth_type="basic", username="USERNAME", password="PASSWORD"
                ),
            )
            ```
        """
        integrated_system_id = self._add_integrated_system(
            integrated_system_credentials,
            name,
            integrated_system_url,
        )

        external_monitor_id = suppress_output(
            self._add_monitor_definitions,
            name,
            metrics,
            schedule,
        )

        # Associate the external monitor with the integrated system
        payload = [
            {
                "op": "add",
                "path": "/parameters",
                "value": {"monitor_definition_ids": [external_monitor_id]},
            },
        ]

        self._wos_client.integrated_systems.update(integrated_system_id, payload)

        return {
            "integrated_system_id": integrated_system_id,
            "monitor_definition_id": external_monitor_id,
        }

    @deprecated(
        reason="'add_observer_instance()' is deprecated and will be removed in a future version. Use 'attach_monitor_instance()' from 'beekeeper-monitors-watsonx' instead.",
        version="1.0.5",
        action="always",
    )
    def add_observer_instance(
        self,
        integrated_system_id: str,
        monitor_definition_id: str,
        subscription_id: str,
    ):
        return self.attach_monitor_instance(
            integrated_system_id=integrated_system_id,
            monitor_definition_id=monitor_definition_id,
            subscription_id=subscription_id,
        )

    @deprecated(
        reason="'add_monitor_instance()' is deprecated and will be removed in a future version. Use 'attach_monitor_instance()' from 'beekeeper-monitors-watsonx' instead.",
        version="1.0.6",
        action="always",
    )
    def add_monitor_instance(
        self,
        integrated_system_id: str,
        monitor_definition_id: str,
        subscription_id: str,
    ):
        return self.attach_monitor_instance(
            integrated_system_id=integrated_system_id,
            monitor_definition_id=monitor_definition_id,
            subscription_id=subscription_id,
        )

    def attach_monitor_instance(
        self,
        integrated_system_id: str,
        monitor_definition_id: str,
        subscription_id: str,
    ):
        """
        Attaches the specified monitor definition to the specified subscription.

        Args:
            integrated_system_id (str): The ID of the integrated system.
            monitor_definition_id (str): The ID of the custom metric monitor instance.
            subscription_id (str): The ID of the subscription to associate the monitor with.

        Example:
            ```python
            wxgov_client.attach_monitor_instance(
                integrated_system_id="019667ca-5687-7838-8d29-4ff70c2b36b0",
                monitor_definition_id="custom_llm_quality",
                subscription_id="0195e95d-03a4-7000-b954-b607db10fe9e",
            )
            ```
        """
        from ibm_watson_openscale.base_classes.watson_open_scale_v2 import Target

        data_marts = self._wos_client.data_marts.list().result.data_marts
        if len(data_marts) == 0:
            raise Exception(
                "No data marts found. Please ensure at least one data mart is available.",
            )

        data_mart_id = data_marts[0].metadata.id
        existing_monitor_instance = self._get_monitor_instance(
            subscription_id,
            monitor_definition_id,
        )

        if existing_monitor_instance is None:
            target = Target(target_type="subscription", target_id=subscription_id)

            parameters = {
                "custom_metrics_provider_id": integrated_system_id,
                "custom_metrics_wait_time": 60,
                "enable_custom_metric_runs": True,
            }

            monitor_instance_details = suppress_output(
                self._wos_client.monitor_instances.create,
                data_mart_id=data_mart_id,
                background_mode=False,
                monitor_definition_id=monitor_definition_id,
                target=target,
                parameters=parameters,
            ).result
        else:
            existing_instance_id = existing_monitor_instance.metadata.id
            monitor_instance_details = self._update_monitor_instance(
                integrated_system_id,
                existing_instance_id,
            )

        return monitor_instance_details

    def publish_metrics(
        self,
        monitor_instance_id: str,
        run_id: str,
        request_records: Dict[str, Union[float, int]],
    ):
        """
        Publishes computed metrics to the specified global monitor instance.

        Args:
            monitor_instance_id (str): The unique ID of the monitor instance.
            run_id (str): The ID of the monitor run that generated the metrics.
            request_records (Dict[str | float | int]): Dict containing the metrics to be published.

        Example:
            ```python
            wxgov_client.publish_metrics(
                monitor_instance_id="01966801-f9ee-7248-a706-41de00a8a998",
                run_id="RUN_ID",
                request_records={"context_quality": 0.914, "sensitivity": 0.85},
            )
            ```
        """
        from ibm_watson_openscale.base_classes.watson_open_scale_v2 import (
            MonitorMeasurementRequest,
            Runs,
        )

        measurement_request = MonitorMeasurementRequest(
            timestamp=datetime.datetime.now(datetime.timezone.utc).strftime(
                "%Y-%m-%dT%H:%M:%S.%fZ",
            ),
            run_id=run_id,
            metrics=[request_records],
        )

        self._wos_client.monitor_instances.add_measurements(
            monitor_instance_id=monitor_instance_id,
            monitor_measurement_request=[measurement_request],
        ).result

        run = Runs(watson_open_scale=self._wos_client)
        patch_payload = []
        patch_payload.append(self._get_patch_request_field("/status/state", "finished"))
        patch_payload.append(
            self._get_patch_request_field(
                "/status/completed_at",
                datetime.datetime.now(datetime.timezone.utc).strftime(
                    "%Y-%m-%dT%H:%M:%S.%fZ",
                ),
            ),
        )

        return run.update(
            monitor_instance_id=monitor_instance_id,
            monitoring_run_id=run_id,
            json_patch_operation=patch_payload,
        ).result

    # ===== Local Custom Metrics =====
    @deprecated(
        reason="'add_local_metric_definition()' is deprecated and will be removed in a future version. Use 'create_local_metric_definition()' from 'beekeeper-monitors-watsonx' instead.",
        version="1.0.6",
        action="always",
    )
    def add_local_metric_definition(
        self,
        name: str,
        metrics: List[WatsonxMetric],
        subscription_id: str,
    ):
        return self.create_local_metric_definition(
            name=name,
            metrics=metrics,
            subscription_id=subscription_id,
        )

    def create_local_metric_definition(
        self,
        name: str,
        metrics: List[WatsonxLocalMetric],
        subscription_id: str,
    ) -> str:
        """
        Creates a custom metric definition to compute metrics at the local (transaction) level for IBM watsonx.governance.

        Args:
            name (str): The name of the custom transaction metric group.
            metrics (List[WatsonxLocalMetric]): A list of metrics to be monitored at the local (transaction) level.
            subscription_id (str): The IBM watsonx.governance subscription ID associated with the metric definition.

        Example:
            ```python
            from beekeeper.monitors.watsonx import WatsonxLocalMetric

            wxgov_client.create_local_metric_definition(
                name="Custom LLM Local Metric",
                subscription_id="019674ca-0c38-745f-8e9b-58546e95174e",
                metrics=[
                    WatsonxLocalMetric(name="context_quality", data_type="double")
                ],
            )
            ```
        """
        from ibm_watson_openscale.base_classes.watson_open_scale_v2 import (
            LocationTableName,
            SparkStruct,
            SparkStructFieldPrimitive,
            Target,
        )

        target = Target(target_id=subscription_id, target_type="subscription")
        data_mart_id = self._get_existing_data_mart()
        metrics = [SparkStructFieldPrimitive(**metric.to_dict()) for metric in metrics]

        schema_fields = [
            SparkStructFieldPrimitive(
                name="scoring_id",
                type="string",
                nullable=False,
            ),
            SparkStructFieldPrimitive(
                name="run_id",
                type="string",
                nullable=True,
            ),
            SparkStructFieldPrimitive(
                name="computed_on",
                type="string",
                nullable=False,
            ),
        ]

        schema_fields.extend(metrics)

        data_schema = SparkStruct(type="struct", fields=schema_fields)

        return self._wos_client.data_sets.add(
            target=target,
            name=name,
            type="custom",
            data_schema=data_schema,
            data_mart_id=data_mart_id,
            location=LocationTableName(
                table_name=name.lower().replace(" ", "_") + "_" + str(uuid.uuid4())[:8],
            ),
            background_mode=False,
        ).result.metadata.id

    def publish_local_metrics(
        self,
        metric_instance_id: str,
        request_records: List[Dict],
    ):
        """
        Publishes computed metrics to the specified transaction record.

        Args:
            metric_instance_id (str): The unique ID of the custom transaction metric.
            request_records (List[Dict]): A list of dictionaries containing the records to be stored.

        Example:
            ```python
            wxgov_client.publish_local_metrics(
                metric_instance_id="0196ad39-1b75-7e77-bddb-cc5393d575c2",
                request_records=[
                    {
                        "scoring_id": "304a9270-44a1-4c4d-bfd4-f756541011f8",
                        "run_id": "RUN_ID",
                        "computed_on": "payload",
                        "context_quality": 0.786,
                    }
                ],
            )
            ```
        """
        return self._wos_client.data_sets.store_records(
            data_set_id=metric_instance_id,
            request_body=request_records,
        ).result

    def list_local_metrics(
        self,
        metric_instance_id: str,
    ):
        """
        Lists records from a custom local metric definition.

        Args:
            metric_instance_id (str): The unique ID of the custom transaction metric.

        Example:
            ```python
            wxgov_client.list_local_metrics(
                metric_instance_id="0196ad47-c505-73c0-9d7b-91c082b697e3"
            )
            ```
        """
        return self._get_dataset_data(metric_instance_id)


@deprecated(
    reason="'WatsonxCustomMetric()' is deprecated and will be removed in a future version. "
    "Use 'WatsonxCustomMetricsManager' instead.",
    version="1.0.6",
    action="always",
)
class WatsonxCustomMetric(WatsonxCustomMetricsManager):
    pass
