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 sientia_do.notifications.models import NotificationLevel
    from sientia_do.temporal.activities.base import BaseActivity
    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

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)

    @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.
        """

        self.logger.debug("Performing input gate...")

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

        filter_output = []

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

        for fil, config in filters.items():
            if fil not in input_filter_functions:
                self.logger.error(f"Filter {fil} not found")
                continue
            try:
                if input_filter_functions[fil](data, config['config']):
                    self.logger.debug(
                        f"Data not passed the input filter {fil}:{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:
                self.logger.debug(f"Input gate result: {path_flag}")
                return path_flag, input_filter_functions['path_confidence'][path_flag], \
                    "Input data with bad quality"

        self.logger.debug("Nothing was filtered by the input gate")
        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.
        """

        self.logger.debug("Performing mlflow response gate...")

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

        filter_output = []

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

        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.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']
                    )
            except Exception as e:
                trace = traceback.format_exc()
                self.notification_handler.build_and_send_notification(
                    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.logger.debug(f"Mlflow response gate result: {path_flag}")
                return path_flag, mlflow_response_filter_functions['path_confidence'][path_flag], \
                    ", ".join(comments)

        self.logger.debug("Nothing was filtered by the mlflow response gate")
        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.
        """

        self.logger.debug("Performing mlflow content gate...")

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

        filter_output = []

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

        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.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()
                    )
            except Exception as e:
                trace = traceback.format_exc()
                self.notification_handler.build_and_send_notification(
                    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.logger.debug(f"Mlflow content gate result: {path_flag}")
                return path_flag, mlflow_content_filter_functions['path_confidence'][path_flag], \
                    "Transformed data not passed the content filter"

        self.logger.debug("Nothing was filtered by the mlflow content gate")
        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.
        """
        self.logger.debug("Formatting prediction...")

        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.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.
        """

        self.logger.debug("Formatting default prediction...")

        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'],
            'comments': [input_data['comment']]
        }).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.
        """
        data = DataFrame(input_data['data'])
        if data.empty:
            return datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        return max(data['timestamp'].values.tolist())