from temporalio import activity, workflow


with workflow.unsafe.imports_passed_through():
    import traceback
    from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
    from sientia_do.notifications.models import NotificationLevel
    from sientia_do.temporal.activities.base import BaseActivity
    from sientia_do.observability.logger import Logger
    from laborious.utils.filters.mlflow_filters import nan_values_filter, api_error_filter
    from typing import Any
    from laborious.utils.filters.conditional_filters import (
        filter_empty_data,
        filter_specific_variables_null_values
    )
    from pandas import DataFrame
    from datetime import datetime
    from laborious import metrics

input_filter_functions = {
    'SPECIFIC_VARIABLES_NULL_VALUES': filter_specific_variables_null_values,
    'EMPTY_DATA': filter_empty_data,
    'path_confidence': {
        'STOP': -1,
        'CONTINUE': 2,
        'REPEAT': -1
    }
}

mlflow_response_filter_functions = {
    'API_ERROR': api_error_filter,
    'path_confidence': {
        'STOP': -1,
        'CONTINUE': 10,
        'REPEAT': -1
    },
}

mlflow_content_filter_functions = {
    'NAN_VALUES': nan_values_filter,
    'path_confidence': {
        'STOP': -1,
        'CONTINUE': 18,
        'REPEAT': -1
    }
}


class Gates(BaseActivity):
    def __init__(self, logger: Logger, notification_handler: NotificationHandler):
        BaseActivity.__init__(
            self, logger, notification_handler, set_error_counter=True)

    @activity.defn(name="input_gate")
    async def input_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
        """
        Filters the data based on the filters. The return value is a tuple with the first element
        being the policy and the second element being the confidence status.
        Args:
            - input_data (dict): The input data. Contains:
                - filters (dict): The filters to apply.
                The key is the filter name and the value is the filter configuration.
                - data (dict[str, Any]): The data to filter.
                - path_priority (list[str]): The path priority.
        Returns:
            tuple[str | None, int, str]: (policy, confidence, comments) based in priority
                list and filter configuration and functions.
        """
        metadata = input_data['metadata']

        self.info("Performing input gate...", metadata)

        self.debug(f"Input data: {input_data}", metadata)

        filters = input_data['filters']
        data = DataFrame(input_data['data'])
        path_priority = input_data['path_priority']

        filter_output = []

        self.debug(f"Input data:\n {data}", metadata)
        self.debug(f"Filters: {filters}", metadata)

        for fil, config in filters.items():
            if fil not in input_filter_functions:
                self.error(f"Filter {fil} not found", metadata)
                continue
            try:
                if input_filter_functions[fil](data, config['config']):
                    self.debug(
                        f"Data not passed the input filter {fil}:{config}", metadata)
                    filter_output.append(config['policy'])
            except Exception as e:
                trace = traceback.format_exc()
                self.send_notification(
                    metadata=metadata,
                    notification_id=f"INTPUT_GATE_ERROR__{fil}",
                    message=f"Error in filter {fil}:{config}: \n {e}",
                    block="input_gate",
                    level=NotificationLevel.ERROR,
                    attachment_content=trace
                )

        for path_flag in path_priority:
            if path_flag in filter_output:
                self.info(f"Input gate result: {path_flag}", metadata)
                return path_flag, input_filter_functions['path_confidence'][path_flag], \
                    "Input data with bad quality"

        self.info("Nothing was filtered by the input gate", metadata)
        return None, 0, ""

    @activity.defn(name="mlflow_response_gate")
    async def mlflow_response_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
        """
        Filters the data based on the mlflow response filters.
        The return value is a tuple with the first element
        being the policy and the second element being the confidence status.
        Args:
            - input_data (dict): The input data. Contains:
                - filters (dict): The filter configuration to apply.
                - data (dict[str, Any]): The data to filter.
                - path_priority (list[str]): The path priority list.
                - type (str): The type of the gate.
        Returns:
            tuple[str | None, int, str]: (policy, confidence, comments) based in priority list
            and filter configuration and functions.
        """

        metadata = input_data['metadata']
        self.info("Performing mlflow response gate...", metadata)

        filters = input_data['filters']
        data = input_data['data']
        gate_type = input_data['type']
        path_priority = input_data['path_priority']

        filter_output = []

        self.debug(f"Input data:\n {data}", metadata)
        self.debug(f"Filters: {filters}", metadata)

        comments = []
        for fil, config in filters.items():
            if fil not in mlflow_response_filter_functions:
                continue
            try:
                if mlflow_response_filter_functions[fil](data, config):
                    filter_output.append(config['policy'])
                    comments.append(data['content']['message'])
                    self.send_notification(
                        metadata=metadata,
                        notification_id=f"{gate_type.upper()}_GATE_RESPONSE_FILTER__{fil}",
                        message=data['content']['message'],
                        block="mlflow_gate",
                        level=NotificationLevel.ERROR,
                        attachment_content=data['content']['traceback']
                    )
            except Exception as e:
                trace = traceback.format_exc()
                self.send_notification(
                    metadata=metadata,
                    notification_id=f"MLFLOW_GATE_RESPONSE_FILTER__{fil}",
                    message=f"Error in filter {fil}:{config}: \n {e}",
                    block="mlflow_gate",
                    level=NotificationLevel.ERROR,
                    attachment_content=trace
                )

        for path_flag in path_priority:
            if path_flag in filter_output:
                self.info(
                    f"Mlflow response gate result: {path_flag}", metadata)
                return path_flag, mlflow_response_filter_functions['path_confidence'][path_flag], \
                    ", ".join(comments)

        self.info("Nothing was filtered by the mlflow response gate", metadata)
        return None, 0, ""

    @activity.defn(name="mlflow_content_gate")
    async def mlflow_content_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
        """
        Filters the data based on the mlflow content filters.
        The return value is a tuple with the first element
        being the policy and the second element being the confidence status.
        Args:
            - input_data (dict): The input data. Contains:
                - filters (dict): The filter configuration to apply.
                - data (dict[str, Any]): The data to filter.
                - path_priority (list[str]): The path priority list.
                - type (str): The type of the gate.
        Returns:
            tuple[str | None, int, str]: (policy, confidence, comments) based in priority
                list and filter configuration and functions.
        """

        metadata = input_data['metadata']
        self.info("Performing mlflow content gate...", metadata)

        filters = input_data['filters']
        data = DataFrame(input_data['data'])
        gate_type = input_data['type']
        path_priority = input_data['path_priority']

        filter_output = []

        self.debug(f"Input data:\n {data}", metadata)
        self.debug(f"Filters: {filters}", metadata)

        for fil, config in filters.items():
            if fil not in mlflow_content_filter_functions:
                continue
            try:
                if mlflow_content_filter_functions[fil](data, config):
                    filter_output.append(config['policy'])
                    self.send_notification(
                        metadata=metadata,
                        notification_id=f"{gate_type.upper()}_GATE_CONTENT_FILTER__{fil}",
                        message=f"Data not passed the content filter {fil}:{config}",
                        block="mlflow_gate",
                        level=NotificationLevel.WARNING,
                        attachment_content=data.to_string()
                    )
            except Exception as e:
                trace = traceback.format_exc()
                self.send_notification(
                    metadata=metadata,
                    notification_id=f"MLFLOW_GATE_CONTENT_FILTER__{fil}",
                    message=f"Error in filter {fil}:{config}: \n {e}",
                    block="mlflow_gate",
                    level=NotificationLevel.ERROR,
                    attachment_content=trace
                )

        for path_flag in path_priority:
            if path_flag in filter_output:
                self.info(
                    f"Mlflow content gate result: {path_flag}", metadata)
                return path_flag, mlflow_content_filter_functions['path_confidence'][path_flag], \
                    "Transformed data not passed the content filter"

        self.info("Nothing was filtered by the mlflow content gate", metadata)
        return None, 0, ""

    @activity.defn(name="format_prediction")
    async def format_prediction(self, input_data: dict[str, Any]) -> dict[Any, Any]:
        """
        Formats the prediction data.
        Args:
            - input_data (dict): The input data. Contains:
                - data (dict[str, Any]): The data to format.
                - timestamp (str): The timestamp of the data.
                - model_id (str): The id of the model.
                - prediction_confidence (float): The confidence of the prediction.
        Returns:
            dict: The formatted data.
        """
        metadata = input_data['metadata']
        self.info("Formatting prediction...", metadata)

        data = DataFrame(input_data['data'])
        data['timestamp'] = input_data['timestamp']
        data['model_id'] = input_data['model_id']
        data['prediction_confidence'] = input_data['prediction_confidence']
        data['prediction_status'] = 'Good'
        data['comments'] = ""
        data = data.sort_values(by='timestamp')

        self.info(f"Prediction formatted: {data.size} rows", metadata)

        return data.to_dict()

    @activity.defn(name="format_default_prediction")
    async def format_default_prediction(self, input_data: dict[str, Any]) -> dict[Any, Any]:
        """
        Creates and formats the default prediction data, with zero value in prediction,
        and usefull information in the other fields.

        Args:
            - input_data (dict): The input data. Contains:
                - timestamp (str): The timestamp of the data.
                - model_id (str): The id of the model.
                - prediction_confidence (float): The confidence of the prediction.
                - comment (str): The comment of the prediction.
        Returns:
            dict: The formatted data.
        """

        metadata = input_data['metadata']
        self.debug("Formatting default prediction...", metadata)

        data = DataFrame({
            'prediction': [0],
            'response_time': [0],
            'timestamp': [input_data['timestamp']],
            'model_id': [input_data['model_id']],
            'prediction_confidence': [input_data['prediction_confidence']],
            'prediction_status': ['Bad'],
            'comments': [input_data['comment']]
        })

        self.info(f"Default prediction formatted: {data.size} rows", metadata)
        return data.to_dict()

    @activity.defn(name="get_last_timestamp")
    async def get_last_timestamp(self, input_data: dict[str, Any]) -> str:
        """
        Gets the last timestamp of the data.
        Args:
            - input_data (dict): The input data. Contains:
                - data (dict[str, Any]): The data to get the last timestamp from.
        Returns:
            str: The last timestamp of the data.
        """
        metadata = input_data['metadata']

        self.info("Getting last timestamp...", metadata)

        data = DataFrame(input_data['data'])

        if data.empty:
            return datetime.now().strftime('%Y-%m-%d %H:%M:%S')

        self.info(
            f"Last timestamp: {max(data['timestamp'].values.tolist())}", metadata)

        return max(data['timestamp'].values.tolist())

    @activity.defn(name="write_metrics")
    async def write_metrics(self, input_data: dict[str, Any]):
        """
        Write metrics to the database.
        input_data:
            metadata: dict[str, Any]
            prediction: dict[str, Any]
        """
        metadata = input_data['metadata']
        prediction = DataFrame(input_data['prediction'])
        prediction_confidence = prediction['prediction_confidence'].values[0]
        response_time = prediction['response_time'].values[0]

        self.info(
            f"Writing metrics for model {metadata['model_name']}", metadata)

        metrics.PREDICTIONS_WRITTEN_COUNT.labels(
            pod_id=self.pod_id,
            model_name=metadata['model_name'],
            pipeline_name=metadata['workflow_name']
        ).inc()

        metrics.PREDICTION_CONFIDENCE_MONITOR.labels(
            pod_id=self.pod_id,
            model_name=metadata['model_name'],
            pipeline_name=metadata['workflow_name']
        ).set(prediction_confidence)

        metrics.PREDICTION_RESPONSE_TIME_MONITOR.labels(
            pod_id=self.pod_id,
            model_name=metadata['model_name'],
            pipeline_name=metadata['workflow_name']
        ).observe(response_time)

        self.info(
            f"Metrics written for model {metadata['model_name']}", metadata)