"""
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.models.train_model_result import TrainModelResult
    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='validate_train_params')
    async def validate_train_params(self, input_data: dict[str, Any]) -> TrainModelParams:
        """
        Validate and convert training parameters from dict to TrainModelParams.

        This activity validates the input training parameters and converts them
        to a TrainModelParams object.

        Args:
            input_data: Training parameters and metadata at the same level
                Required keys:
                    - metadata (dict): Workflow execution metadata
                    - All TrainModelParams fields (experiment_run_id, target_variable, etc.)

        Returns:
            TrainModelParams: Validated and converted training parameters

        Raises:
            ValueError, TypeError, KeyError: If validation fails (after sending notification)

        Example:
            result = await validate_train_params({
                'metadata': {'workflow_id': 'train-123'},
                'experiment_run_id': 456,
                'target_variable': 'price',
                'variable_columns': ['feature1', 'feature2'],
                'train_size': 80,
                # ... other required fields at same level
            })
            # Returns: TrainModelParams(...)
        """
        metadata = input_data.get('metadata', {})

        try:
            self.info('Validating training parameters', metadata)

            # Step 1: Convert input_data to TrainModelParams (validates types and required fields)
            train_params = TrainModelParams.from_dict(input_data)

            # Step 2: Validate business rules (ranges, consistency, etc.)
            train_params.validate_business_rules()

            self.info(
                f'Training parameters validated successfully - '
                f'Target: {train_params.target_variable}, '
                f'Experiment: {train_params.experiment_name}',
                metadata,
            )

            return train_params

        except (ValueError, TypeError, KeyError) as e:
            error_msg = f'Error validating training parameters: {str(e)}'
            trace = traceback.format_exc()

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

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

            # Re-raise exception to stop workflow
            raise

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

        This activity orchestrates the complete ML training pipeline:
        1. Validates input parameters
        2. Trains the model using TrainingRepository
        3. Performs post-training calculations

        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 (TrainModelParams): Training parameters object

        Returns:
            TrainModelResult: Training result with model, metrics, and data

        Raises:
            ValueError: If input validation fails
            Exception: If training fails (after sending notification)

        Example:
            result = await train_model({
                'metadata': {'workflow_id': 'train-123', 'experiment_run_id': 456},
                'uploaded_file': BytesIO(csv_data),
                'train_params': TrainModelParams(...)
            })
            # Returns: TrainModelResult(...)
        """
        metadata = input_data.get('metadata', {})
        uploaded_file = input_data['uploaded_file']
        train_params = input_data['train_params']

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

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

            self.info(
                f'Starting model training for target: {train_params.target_variable}',
                metadata,
            )

            # 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 final_result

        except Exception as e:  # noqa: BLE001
            target = (
                train_params.target_variable
                if hasattr(train_params, 'target_variable')
                else 'unknown'
            )
            error_msg = f'Error training model - Target: {target}, 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)

            # Re-raise exception to stop workflow
            raise