from temporalio import workflow

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

    from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
    from sientia_do.temporal.policies import retry_policy

    from laborious.activities.activities import Activities


@workflow.defn(name='simple_metrics')
class SimpleMetrics:
    @workflow.run
    async def run(self, input_data: dict[str, Any]):
        """
        Execute the simple metrics workflow.
        """
        metadata = {
            'metadata': {
                'model_id': input_data['model_id'],
                'model_name': input_data['model_name'],
                'workflow_name': 'simple_metrics',
                'schedule_name': input_data['schedule_name'],
            }
        }

        model_id = input_data['model_id']
        interval_minutes = input_data['interval_minutes']

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

        query = f"""
            select p."timestamp", p.prediction, ld.value as "target"
            from "{input_data['schema']}"."{input_data['predictions_table_name']}" p
            inner join "{input_data['schema']}"."{input_data['data_table_name']}" ld
                on p."timestamp" = ld."timestamp"
            where
                p.model_id = '{model_id}' and
                p.prediction is not null and
                ld.variable = '{target_name}' and
                ld.value is not null and
                p."timestamp" >= NOW() - INTERVAL '{interval_minutes} minutes'
            order by
                p."timestamp" desc;
        """  # nosec B608 - values come from internal Temporal workflow config, not user input

        target_data = await workflow.execute_activity_method(
            Activities.load_custom_query,
            {
                **metadata,
                'query': query,
                'datetime_columns': ['timestamp'],
                'orient': 'records',
            },
            retry_policy=retry_policy,
            start_to_close_timeout=timedelta(seconds=300),
        )

        if not target_data:
            return

        simple_metrics = await workflow.execute_local_activity_method(
            Activities.calculate_simple_metrics,
            {
                **metadata,
                'model_id': model_id,
                'target_data': target_data,
                'metrics': input_data.get('metrics', ['rmse', 'mse', 'mae', 'r2']),
                'interval_minutes': interval_minutes,
            },
            retry_policy=retry_policy,
            start_to_close_timeout=timedelta(seconds=300),
        )

        if not simple_metrics:
            return

        await workflow.execute_activity_method(
            Activities.export_data_to_postgres,
            {
                **metadata,
                'data': simple_metrics,
                '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),
        )