"""
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 typing import Any

    import pandas as pd
    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.observability.metrics_controller import MetricsController
    from sientia_do.observability.sientia_monitoring import SientiaMonitoring
    from sientia_do.repository.minio_repository import MinioRepository
    from sientia_model.model_repository.mlflow_repository import SientiaMLflowRepository
    from sientia_model.model_repository.plugin_store import PluginStore
    from sientia_model.wrappers.sientia_model import SientiaModel

    from model_manager.metrics import ACTIVITY_EXECUTION_TOTAL, WORKFLOW_EXECUTION_TOTAL
    from model_manager.utils.exceptions import ModelTrainingError
    from model_manager.utils.models.train_model_params import TrainModelParams
    from model_manager.utils.repository.data_manager_repository import DataManagerRepository


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

    This activity extends SientiaMonitoring 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.
    """

    def __init__(
        self,
        mlflow_repository: SientiaMLflowRepository,
        plugin_store: PluginStore,
        minio_repository: MinioRepository,
        logger: Logger,
        notification_handler: NotificationHandler,
        metrics_controller: MetricsController,
    ):
        """
        Initialize Training activity.

        Args:
            logger: Logger instance for observability
            notification_handler: Handler for sending notifications
        """
        SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
        self.data_manager_repository = DataManagerRepository(logger)
        self.mlflow_repository = mlflow_repository
        self.plugin_store = plugin_store
        self.minio_repository = minio_repository

    @activity.defn(name='load_model_metadata')
    async def load_model_metadata(self, input_data: dict[str, Any]) -> dict[str, Any]:
        """
        Load model metadata/schemas from the model store.

        This activity is responsible for fetching model metadata/schemas from the
        model store index and extracting a serializable `model_metadata` dict that
        `TrainModelParams.validate_business_rules()` depends on.

        Args:
            input_data: Workflow input at the same level as `validate_train_params`,
                including at least `model_name` and the fields required by
                `TrainModelParams.from_dict` to build wrapper kwargs.

        Return:
            dict[str, Any]: Updated `input_data` containing `input_data['model_metadata']`.
        """
        metadata = input_data.get('metadata', {})

        try:
            train_params = TrainModelParams.from_dict(input_data)
            model_metadata = self.plugin_store.get_model_index(
                model_name=train_params.model_name,
                metadata=metadata,
            )

            train_params.model_metadata = model_metadata
            return train_params.to_dict()
        except Exception as exc:
            trace = traceback.format_exc()
            await self.send_notification_async(
                metadata=metadata,
                notification_id='LOAD_MODEL_METADATA_ERROR',
                message=f'Error loading model metadata: {str(exc)}',
                block='load_model_metadata',
                level=NotificationLevel.ERROR,
                attachment_content=trace,
            )
            raise
        
    @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)
        """
        metadata = input_data.get('metadata', {})
        try:
            train_params = TrainModelParams.from_dict(input_data)
            
            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()

            await self.send_notification_async(
                metadata=metadata,
                notification_id='VALIDATE_TRAIN_PARAMS_ERROR',
                message=error_msg,
                block='validate_train_params',
                level=NotificationLevel.ERROR,
                attachment_content=trace,
            )
            raise

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

        This activity orchestrates the ML training pipeline:
        1. Validate input parameters.
        2. Prepare data via DataManagerRepository.
        3. Train the model and compute metrics.

        Args:
            input_data: Training configuration containing:
                - metadata (dict): Workflow execution metadata.
                - uploaded_file (BytesIO): Training data already downloaded from MinIO.
                - train_params (TrainModelParams | dict): Training parameters.

        Returns:
            dict: Key `run_name` when training and saving succeed.

        Raises:
            ValueError: If input validation fails.
            Exception: If training fails (after sending notification).
        """
        metadata = input_data.get('metadata')
        train_params = input_data['train_params']

        if isinstance(train_params, dict):
            train_params = TrainModelParams.from_dict(train_params)

        model_trained = False
        model_saved = False

        try:
            # Download training file bytes from MinIO
            train_bytes = await self.minio_repository.download_file(
                object_name=train_params.file_name,
                bucket=train_params.bucket_name,
                metadata=metadata,
            )

            # Download optional validation file bytes from the same bucket
            val_bytes: bytes | None = None
            validation_name = train_params.val_file_name
            if validation_name is not None:
                val_bytes = await self.minio_repository.download_file(
                    object_name=validation_name,
                    bucket=train_params.bucket_name,
                    metadata=metadata,
                )

            train_result = self.data_manager_repository.prepare_training_data(
                train_file_bytes=train_bytes,
                validation_file_bytes=val_bytes,
                params=train_params,
                metadata=metadata,
            )

            wrapper = await self.plugin_store.get_model(
                model_name=train_params.model_name,
                force_download=False,
                opt_params=train_params.opt_params or {},
                model_kwargs=train_params.model_kwargs or {},
                data_model_kwargs=train_params.data_model_kwargs or {},
                metadata=metadata
            )

            train_data = train_result.train_data
            val_data = train_result.val_data

            wrapper.train(
                train_data=train_data,
                val_data=val_data,
                target=train_params.target_variable,
            )

            # Generate predictions using the trained wrapper
            transformed_train, _ = wrapper.transform(train_data)
            transformed_val, _ = wrapper.transform(val_data)

            y_train_pred_df, _ = wrapper.predict({}, transformed_train)
            y_val_pred_df, _ = wrapper.predict({}, transformed_val)

            y_train_pred_df.sort_index(inplace=True, ascending=False)
            y_val_pred_df.sort_index(inplace=True, ascending=False)

            train_result.y_train_pred = y_train_pred_df
            train_result.y_pred = y_val_pred_df

            train_result = self.data_manager_repository.compute_regression_metrics(
                train_result,
            )

            model_trained = True

            async with self.mlflow_repository.start_run(
                model_name=train_params.model_name,
                run_name=None,
                experiment_name=f'{train_params.model_name}_experiment',
                tags=None,
                metadata=metadata,
            ) as run_info:
                wrapper.store_model(name=train_params.model_name)

            model_saved = True

            return {
                'run_name': run_info.run_name,
                'run_id': run_info.run_id,
            }
        except Exception as e:  # noqa: BLE001
            metrics_status = 'error'

            error_msg = (
                'Error training model - '
                f'model_trained={model_trained}, model_saved={model_saved}, '
                f'error: {str(e)}'
            )

            trace = traceback.format_exc()

            await self.send_notification_async(
                metadata=metadata or {},
                notification_id='TRAIN_MODEL_ERROR',
                message=error_msg,
                block='train_model',
                level=NotificationLevel.ERROR,
                attachment_content=trace,
            )

            raise ModelTrainingError(
                model_trained=model_trained,
                model_saved=model_saved,
            ) from e

    @activity.defn(name='cleanup_resources')
    async def cleanup_resources(self, input_data: dict[str, Any]) -> None:
        """
        Cleanup temporary resources created during training.

        Args:
            input_data: Cleanup configuration containing:
                - metadata (dict): Workflow execution metadata.
                - bucket_name (str): MinIO bucket of the uploaded file.
                - file_name (str): MinIO object key to delete.

        Raises:
            Exception: If cleanup fails (after sending notification).
        """
        metadata = input_data.get('metadata', {})
        bucket_name = input_data.get('bucket_name', '')
        file_name = input_data.get('file_name', '')
        metrics_status = 'success'

        try:
            await self.minio_repository.delete_file(
                object_name=file_name,
                bucket=bucket_name,
                metadata=metadata,
            )
        except Exception as e:  # noqa: BLE001
            error_msg = (
                'Error cleaning up resources - '
                f'File: {bucket_name}/{file_name}, Error: {str(e)}'
            )

            trace = traceback.format_exc()

            self.send_notification(
                metadata=metadata,
                notification_id='CLEANUP_RESOURCES_ERROR',
                message=error_msg,
                block='cleanup_resources',
                level=NotificationLevel.ERROR,
                attachment_content=trace,
            )

            raise