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 model_manager.activities.activities import Activities


@workflow.defn(name='minimal_retrain')
class MinimalRetrain:
    """
    Automated model retraining workflow for the Model Manager system.

    This workflow implements a complete model retraining pipeline that loads
    training data, executes model retraining, updates production models,
    and maintains comprehensive audit trails. It's designed for automated
    model lifecycle management with minimal manual intervention.

    The workflow provides a robust retraining process with:
    - Automated data loading from configured data sources
    - MLFlow model retraining with quality validation
    - Production model updates with version control
    - Comprehensive reporting and audit trail maintenance
    - Error handling and notification integration
    """

    @workflow.run
    async def run(self, input_data: dict[str, Any]):
        """
        Execute the automated model retraining workflow.

        This method orchestrates the complete model retraining process by:
        1. Loading training data using the provided custom SQL query
        2. Executing MLFlow model retraining with the loaded data
        3. Updating production models with newly trained versions
        4. Persisting comprehensive retraining reports to database

        The method implements comprehensive error handling and ensures all
        required parameters are properly configured before proceeding.

        Args:
            input_data: Complete configuration for the retraining workflow
                Required keys:
                    - schedule_name (str): Schedule identifier for the retraining
                    - model_name (str): Name of the ML model to retrain
                    - model_id (int): Unique identifier for the model version
                    - query (str): SQL query for training data loading
                    - schema (str, optional): Database schema for report storage
                    - table_name (str, optional): Target table for retraining reports
                    - datetime_columns (list[str], optional): Columns to treat as datetime

        Returns:
            None: The workflow completes successfully when all steps finish

        Raises:
            Exception: If any required parameters are missing or if the workflow fails
                during data loading, retraining, or model update operations
        """

        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'],
                'datetime_columns': input_data.get('datetime_columns', []),
            },
            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),
        )