from temporalio import workflow, activity

with workflow.unsafe.imports_passed_through():
    from sientia_do.notifications.models import NotificationLevel
    from pandas import DataFrame
    from scouter.activities.base import BaseActivity
    from scouter.utils.quality.filters import null_values_filter, out_of_bounds_filter
    from typing import Any
    import traceback

quality_gate_filters = {
    'NULL_VALUES_FILTER': null_values_filter,
    'OUT_OF_BOUNDS_FILTER': out_of_bounds_filter
}


class Gates(BaseActivity):
    @activity.defn(name="data_quality_gate")
    async def data_quality_gate(self, input_data: dict[str, Any]) -> dict[str, Any]:
        """
        Data quality gate activity. for each selected filter,
        extracts filtered data, discards or keeps filtered data
        based on the filter.

        Args:
            input_data (dict[str, Any]): The data to validate. Contains:
                filters (dict[str, str]): The filters to apply. In format:
                    {filter_name: policy}.
                    filter_name: The name of the filter.
                    policy: The policy to apply. Can be "DISCARD" or "KEEP".
                data (dict[str, Any]): The data to validate.

        Returns:
            dict[str, Any]: The data validated.
        """

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

        for filter_name, policy in filters.items():
            if filter_name not in quality_gate_filters:
                self.logger.warning(f"Filter {filter_name} not found")
                continue

            try:
                filtered_data = quality_gate_filters[filter_name](data)

            except Exception as e:
                self.notification_handler.build_and_send_notification(
                    notification_id="DATA_QUALITY_GATE_ISSUES",
                    message=f"Error applying filter {filter_name}: {e}",
                    block="data_quality_gate",
                    level=NotificationLevel.ERROR,
                    attachment_content=traceback.format_exc()
                )

            else:
                if filtered_data.empty:
                    continue

                message = f"{len(filtered_data)} rows has quality issues: {filter_name}: {policy}"
                attachment = filtered_data.to_string()

                self.notification_handler.build_and_send_notification(
                    notification_id="DATA_QUALITY_GATE_ISSUES",
                    message=message,
                    block="data_quality_gate",
                    level=NotificationLevel.WARNING,
                    attachment_content=attachment
                )

                if policy == "DISCARD":
                    data = data[not data.isin(filtered_data).all(axis=1)]

        return data.to_dict()