from temporalio import activity, workflow


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

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):
        super().__init__(logger, notification_handler)

    @activity.defn(name="input_gate")
    async def input_gate(self, input_data: dict[str, Any]) -> tuple[str, int]:
        """
        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.
                data (dict[str, Any]): The data to filter.
                path_priority (list[str]): The path priority.
        Returns:
            tuple[str, int]: (policy, confidence) based in priority list and filter configuration and functions.
        """
        filters = input_data['filters']
        data = DataFrame(input_data['data'])
        path_priority = input_data['path_priority']

        filter_output = []
        for fil, config in filters.items():
            try:
                if input_filter_functions[fil](data, config):
                    filter_output.append(config['POLICY'])
            except Exception as e:
                trace = traceback.format_exc()
                self.notification_handler.build_and_send_notification(
                    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:
                return path_flag, input_filter_functions['path_confidence'][path_flag]

        return None, 0

    @activity.defn(name="mlflow_response_gate")
    async def mlflow_response_gate(self, input_data: dict[str, Any]) -> tuple[str, int]:
        """
        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, int]: (policy, confidence) based in priority list and filter configuration and functions.
        """
        filters = input_data['filters']
        data = input_data['data']
        gate_type = input_data['type']
        path_priority = input_data['path_priority']

        filter_output = []
        for fil, config in filters.items():
            if mlflow_response_filter_functions[fil](data, config):
                filter_output.append(config['POLICY'])
                self.notification_handler.build_and_send_notification(
                    notification_id=f"{gate_type.upper()}_GATE_RESPONSE_FILTER__{fil}",
                    message=data['content']['message'],
                    block="mlflow_gate",
                    level=NotificationLevel.WARNING,
                    attachment_content=data['content']['traceback']
                )

        for path_flag in path_priority:
            if path_flag in filter_output:
                return path_flag, mlflow_response_filter_functions['path_confidence'][path_flag]

        return None, 0

    @activity.defn(name="mlflow_content_gate")
    async def mlflow_content_gate(self, input_data: dict[str, Any]) -> tuple[str, int]:
        """
        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, int]: (policy, confidence) based in priority list and filter configuration and functions.
        """

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

        filter_output = []

        for fil, config in filters.items():
            if mlflow_content_filter_functions[fil](data, config):
                filter_output.append(config['POLICY'])
                self.notification_handler.build_and_send_notification(
                    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()
                )

        for path_flag in path_priority:
            if path_flag in filter_output:
                return path_flag, mlflow_content_filter_functions['path_confidence'][path_flag]

        return None, 0

    @activity.defn(name="format_prediction")
    async def format_prediction(self, input_data: dict[str, Any]) -> dict[str, 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.
        """
        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['comment'] = ""
        data.sort_values(by='timestamp', inplace=True)

        return data.to_dict()

    @activity.defn(name="format_default_prediction")
    async def format_default_prediction(self, input_data: dict[str, Any]) -> dict[str, 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.
        """

        return 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'],
            'comment': [input_data['comment']]
        }).to_dict()