"""
Training activities for ML model training operations.

This module provides activities for training machine learning models.
The activity extends BaseActivity and receives pre-downloaded files
to return success/failure status without raising exceptions.
"""

from temporalio import activity, workflow

with workflow.unsafe.imports_passed_through():
    import traceback
    from io import BytesIO
    from typing import Any

    from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
    from sientia_do.notifications.models import NotificationLevel
    from sientia_do.observability.logger import Logger
    from sientia_do.temporal.activities.base import BaseActivity

    from model_manager.utils.models.train_model_params import TrainModelParams
    from model_manager.utils.repository.training_repository import TrainingRepository


class Training(BaseActivity):
    """
    Activity for ML model training operations.

    This activity extends BaseActivity and handles machine learning model
    training with comprehensive error handling. It receives pre-downloaded
    files from the workflow and returns success/failure status without
    raising exceptions.

    Attributes:
        logger (Logger): Logger instance for observability (inherited from BaseActivity)
        notification_handler (NotificationHandler): Handler for sending notifications (inherited)
    """

    def __init__(
        self,
        logger: Logger,
        notification_handler: NotificationHandler,
    ):
        """
        Initialize Training activity.

        Args:
            logger: Logger instance for observability
            notification_handler: Handler for sending notifications
        """
        BaseActivity.__init__(self, logger, notification_handler, set_error_counter=True)
        self.training_repository = TrainingRepository(logger)

    @activity.defn(name='train_model')
    async def train_model(self, input_data: dict[str, Any]) -> dict[str, Any]:
        """
        Train a machine learning model with comprehensive error handling.

        This activity orchestrates the complete ML training pipeline:
        1. Validates input parameters
        2. Trains the model using TrainingRepository
        3. Performs post-training calculations
        4. Returns success/failure status with results or error message

        The activity does NOT raise exceptions on failure - it catches all errors,
        sends notifications, and returns a failure status. This allows the workflow
        to handle the error gracefully and update the database accordingly.

        Args:
            input_data: Configuration for model training operation
                Required keys:
                    - metadata (dict): Workflow execution metadata
                    - uploaded_file (BytesIO): Training data file (already downloaded from MinIO)
                    - train_params (dict): Training parameters (converted to TrainModelParams)

        Returns:
            dict: Training result with the following structure:
                {
                    'success': bool,  # True if training succeeded, False otherwise
                    'result': TrainModelResult | None,  # Training result if success=True
                    'error_message': str | None  # Error message if success=False
                }

        Example:
            # Successful training
            result = await train_model({
                'metadata': {'workflow_id': 'train-123', 'experiment_run_id': 456},
                'uploaded_file': BytesIO(csv_data),
                'train_params': {
                    'experiment_run_id': 456,
                    'target_variable': 'price',
                    'variable_columns': ['feature1', 'feature2'],
                    'train_size': 80,
                    'shuffle': True,
                    'use_scaler': True,
                    # ... other TrainModelParams fields
                }
            })
            # Returns: {'success': True, 'result': TrainModelResult(...), 'error_message': None}

            # Failed training
            # Returns: {'success': False, 'result': None, 'error_message': 'Error details...'}
        """
        metadata = input_data.get('metadata', {})
        uploaded_file = input_data['uploaded_file']
        train_params_dict = input_data['train_params']

        try:
            self.info(
                f'Starting model training for target: {train_params_dict.get("target_variable")}',
                metadata,
            )

            # Convert dict to TrainModelParams
            train_params = TrainModelParams.from_dict(train_params_dict)

            # Validate uploaded_file is BytesIO
            if not isinstance(uploaded_file, BytesIO):
                raise ValueError(f'uploaded_file must be BytesIO, got {type(uploaded_file)}')

            # Step 1: Train the model
            self.info('Training model with TrainingRepository', metadata)
            train_result = self.training_repository.train(uploaded_file, train_params)

            # Step 2: Perform post-training calculations
            self.info('Performing post-training calculations', metadata)
            final_result = self.training_repository.after_train_calculation(
                train_params, train_result
            )

            self.info(
                f'Model training completed successfully - '
                f'MSE: {final_result.mse_val}, MAE: {final_result.mae_val}, R²: {final_result.r2_val}',
                metadata,
            )

            return {
                'success': True,
                'result': final_result,
                'error_message': None,
            }

        except Exception as e:  # noqa: BLE001
            error_msg = f'Error training model - Target: {train_params_dict.get("target_variable", "unknown")}, Error: {str(e)}'
            trace = traceback.format_exc()

            # Send notification (MongoDB)
            self.send_notification(
                metadata=metadata,
                notification_id='TRAIN_MODEL_ERROR',
                message=error_msg,
                block='train_model',
                level=NotificationLevel.ERROR,
                attachment_content=trace,
            )

            # Log error with metadata
            self.error(trace, metadata=metadata)

            # Return failure result (do NOT raise exception)
            # This allows workflow to update database with error status
            return {
                'success': False,
                'result': None,
                'error_message': str(e),
            }