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
}

transform_filter_functions = {
    'response_filter': {
        'API_ERROR': api_error_filter,
    },
    'content_filter': {
        'NAN_VALUES': nan_values_filter,
    }
}


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.
        Returns:
            tuple[str, int]: ('stop', -1) if some filter policy is 'stop', ('continue', 2)
            if no filter policy is 'stop' and some filter policy is 'continue',
            None if no filter is applied.
        """
        filters = input_data['filters']
        data = DataFrame(input_data['data'])

        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
                )

        if 'stop' in filter_output:
            return 'stop', -1
        elif 'continue' in filter_output:
            return 'continue', 2

        return None, 0

    @activity.defn(name="mlflow_gate")
    async def mlflow_gate(self, input_data: dict[str, Any]) -> tuple[str, int]:

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

        filter_output = []
        for fil, config in filters.items():
            if transform_filter_functions['response_filter'][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']
                )

        if 'stop' in filter_output:
            return 'stop', -1
        elif 'continue' in filter_output:
            return 'continue', 10

        if gate_type == 'predict':
            return None, 0

        data = DataFrame(data['content'])

        for fil, config in filters.items():
            if transform_filter_functions['content_filter'][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()
                )

        if 'stop' in filter_output:
            return 'stop', -1
        elif 'continue' in filter_output:
            return 'continue', 18

        return None, 0

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