from temporalio import workflow

with workflow.unsafe.imports_passed_through():
    from datetime import timedelta
    from typing import Any

    from sientia_do.temporal.policies import retry_policy

    from laborious.activities.activities import Activities
    from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ


@workflow.defn(name='drift')
class Drift:
    @workflow.run
    async def run(self, input_data: dict[str, Any]):
        """
        Execute the drift workflow.

        This method orchestrates the complete drift process by:
        1. Loading data using the provided custom SQL query
        2. Preparing prediction configuration and filters
        3. Delegating to the PredictionProcess workflow for ML operations
        """

        metadata = {
            'metadata': {
                'schedule_name': input_data['schedule_name'],
                'model_name': input_data['model_name'],
                'model_id': input_data['model_id'],
                'workflow_name': 'drift',
            }
        }

        model_config = input_data['model_config']
        target_name = model_config['target']

        gathering_query = f"""
            SELECT *
            FROM {input_data['schema']}.{input_data['source_table_name']}
            WHERE
                model_id = {input_data['model_id']} AND
                timestamp > NOW() - INTERVAL '{input_data['interval']} minutes'
            ORDER BY timestamp ASC
        """

        target_data_handler = workflow.start_local_activity_method(
            Activities.load_custom_query,
            {
                **metadata,
                'query': gathering_query,
                'datetime_columns': ['timestamp', 'created_at'],
            },
            retry_policy=retry_policy,
            start_to_close_timeout=timedelta(seconds=300),
        )

        reference_data_handler = workflow.start_local_activity_method(
            Activities.get_reference_data,
            {
                **metadata,
                'model_name': input_data['model_name']
            },
            retry_policy=retry_policy,
            start_to_close_timeout=timedelta(seconds=300),
        )

        target_data = await target_data_handler
        reference_data = await reference_data_handler

        if not target_data:
            return

        drift_data = await workflow.execute_local_activity_method(
            Activities.calculate_drift,
            {
                **metadata,
                'target_data': target_data,
                'reference_data': reference_data,
                'model_name': input_data['model_name'],
                'model_id': input_data['model_id'],
                'target_name': target_name,
                'drift_metrics': input_data.get('drift_metrics',
                    ['kolmogorov_smirnov', 'jensen_shannon', 'wasserstein']),
                'chunk_period': input_data.get('chunk_period', 'min'),
            },
            retry_policy=retry_policy,
            start_to_close_timeout=timedelta(seconds=300),
        )

        if drift_data:
            await workflow.execute_activity_method(
                Activities.export_data_to_postgres,
                {
                    **metadata,
                    'data': drift_data,
                    'schema': input_data['schema'],
                    'table_name': input_data['target_table_name'],
                    'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ},
                },
                retry_policy=retry_policy,
                start_to_close_timeout=timedelta(seconds=300),
            )