from temporalio import workflow

with workflow.unsafe.imports_passed_through():
    from laborious.activities.activities import Activities
    from typing import Any
    from sientia_do.temporal.utils.policies import retry_policy
    from datetime import timedelta


@workflow.defn(name="predictions_batch")
class PredictionsBatch():
    @workflow.run
    async def run(self, input_data: dict[str, Any]):
        """
        This workflow runs a batch of predictions based on the input data.

        The workflow executes in two main steps:
        1. Prepares the activity with schedule and model information
        2. Loads data using a custom query and executes the prediction process

        Args:
            - input_data (dict[str, Any]): The input data for the workflow.
            Contains the following keys:
                - schedule_name (str): The name of the schedule.
                - model_name (str): The name of the model.
                - model_id (int): The id of the model.
                - query (str): The SQL query to be executed to load data.
                - schema (dict, optional): The schema definition for the data.
                - table_name (str, optional): The name of the table to process.
                - input_filters (dict, optional): Filters to be applied during prediction.
                - mlflow_transform_filters (dict, optional): Filters to be applied
                during prediction.
                - mlflow_predict_filters (dict, optional): Filters to be applied during prediction.
                - model_retention (int, optional): The model retention period in minutes.
                - path_priority (list[str]): The path priority.
        Returns:
            None

        Raises:
            Exception: If any of the required parameters are missing or if the workflow fails.
        """

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

        data = await workflow.execute_local_activity_method(
            Activities.load_custom_query,
            {
                **metadata,
                'query': input_data['query'],
            },
            retry_policy=retry_policy,
            start_to_close_timeout=timedelta(seconds=60)
        )

        # Prepare input for prediction_process workflow
        prediction_input = {
            'metadata': metadata,
            'data': data,
            'schema': input_data['schema'],
            'table_name': input_data['table_name'],
            'model_id': input_data['model_id'],
            'model_name': input_data['model_name'],
            'input_filters': input_data.get('input_filters', {
                'EMPTY_DATA': {
                    'POLICY': 'STOP'
                }
            }),
            'mlflow_transform_filters': input_data.get('mlflow_transform_filters', {
                'API_ERROR': {
                    'POLICY': 'STOP'
                }
            }),
            'mlflow_predict_filters': input_data.get('mlflow_predict_filters', {
                'API_ERROR': {
                    'POLICY': 'STOP'
                }
            }),
            'model_retention': input_data.get('model_retention', 60),
            'path_priority': input_data.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT']),
            'opc_output_config': input_data.get('opc_output_config', {})
        }

        await workflow.execute_child_workflow(
            'prediction_process', prediction_input)