from temporalio import workflow

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


@workflow.defn(name="minimal_retrain")
class MinimalRetrain():
    @workflow.run
    async def run(self, input_data: dict[str, Any]):
        """
        This workflow runs a minimal retrain of a model.

        The workflow executes in four steps:
        1. Loads the data from the database
        2. Formats the data and perform the retrain
        3. Updates the production model
        4. Saves a model

        Args:
            - input_data (dict[str, Any]): The input data for the workflow.
                - 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 to store the report.
                - table_name (str, optional): The name of the table to store report.

        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': 'minimal_retrain'
            }
        }

        model_name = input_data['model_name']

        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)
        )

        experiment_response = await workflow.execute_activity_method(
            Activities.retrain_model,
            {
                **metadata,
                'data': data,
                'model_name': model_name
            },
            retry_policy=retry_policy,
            start_to_close_timeout=timedelta(seconds=60)
        )

        report = await workflow.execute_activity_method(
            Activities.update_production_model,
            {
                **metadata,
                'model_name': model_name,
                'model_id': input_data['model_id'],
                **experiment_response
            },
            retry_policy=retry_policy,
            start_to_close_timeout=timedelta(seconds=60)
        )

        await workflow.execute_activity_method(
            Activities.export_data_to_postgres,
            {
                **metadata,
                'data': report,
                'schema': input_data['schema'],
                'table_name': input_data['table_name']
            },
            retry_policy=retry_policy,
            start_to_close_timeout=timedelta(seconds=60)
        )