"""
Experiment tracking activities for managing ML experiment lifecycle.

This module provides activities for tracking and updating experiment run status
in the PostgreSQL database, extending the synchronous Postgres client with specialized
methods for experiment management.
"""

import enum

from temporalio import activity, workflow

with workflow.unsafe.imports_passed_through():
    import traceback
    from collections.abc import Mapping
    from datetime import UTC, datetime
    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.observability.metrics_controller import MetricsController
    from sientia_do.temporal.activities.postgres_sync import Postgres
    from sqlalchemy import text


class UpdateType(enum.StrEnum):
    """Types of experiment run updates."""

    STATUS = 'status'
    STATUS_WITH_ERROR = 'status_with_error'
    MODEL_SAVED = 'model_saved'


class ExperimentTracking(Postgres):
    """
    Activity for tracking ML experiment lifecycle and status updates.

    This activity extends the Postgres activity to provide specialized methods
    for managing experiment runs, including status updates, error tracking, and
    model registration. It maintains the experiment lifecycle from initialization
    through training, model saving, and cleanup.

    The activity uses a SQLAlchemy engine / connection pool and adds experiment-specific
    operations with proper error handling and notifications.
    """

    def __init__(
        self,
        host: str,
        port: int,
        user: str,
        password: str,
        dbname: str,
        min_connections: int,
        max_connections: int,
        logger: Logger,
        notification_handler: NotificationHandler,
        metrics_controller: MetricsController,
    ):
        """
        Initialize ExperimentTracking activity with database configuration.

        Args:
            host: PostgreSQL server hostname
            port: PostgreSQL server port
            user: Database user
            password: Database password
            dbname: Database name
            min_connections: Minimum connections in pool
            max_connections: Maximum connections in pool
            logger: Logger instance for observability
            notification_handler: Notification handler for alerts
            metrics_controller: Metrics controller for observability

        Raises:
            ConnectionError: If database connection cannot be established
        """
        Postgres.__init__(
            self,
            host=host,
            port=port,
            user=user,
            password=password,
            dbname=dbname,
            min_connections=min_connections,
            max_connections=max_connections,
            logger=logger,
            notification_handler=notification_handler,
            metrics_controller=metrics_controller,
        )

        self.info(f'Postgres client initialized at {host}:{port}')

    def __del__(self):
        """
        Destructor to safely handle cleanup during garbage collection.

        This prevents AttributeError when used in multiple inheritance scenarios
        where the parent Postgres.__del__ might be called on objects without
        the engine attribute.
        """
        # Only call parent __del__ if engine attribute exists
        if hasattr(self, 'engine'):
            try:
                if hasattr(super(), '__del__'):
                    super().__del__()
            except Exception:  # noqa: S110, BLE001
                # Silently ignore errors during garbage collection
                pass

    def _execute_update(self, query: str, params: Mapping[str, Any]) -> dict[str, Any]:
        """
        Execute an UPDATE SQL statement.

        Args:
            query: Parameterized SQL string to execute.
            params: Mapping of parameters for the SQL query.

        Returns:
            dict: A dictionary containing the affected row count: {'rowcount': int}.
        """

        with self.engine.begin() as connection:
            result = connection.execute(text(query), params)
            return {'rowcount': result.rowcount}

    def _build_status_update_query(
        self, status: str | None, experiment_run_id: int
    ) -> tuple[str, dict[str, Any]]:
        """Build SQL query for simple status update."""
        if not isinstance(status, str) or not status:
            raise ValueError('status is required for STATUS update type')

        sql_query = """
            UPDATE experiment_run
            SET status = :status, updated_at = :updated_at
            WHERE id = :experiment_run_id
        """

        query_params = {
            'status': status,
            'updated_at': datetime.now(UTC),
            'experiment_run_id': experiment_run_id,
        }

        return sql_query, query_params

    def _build_status_with_error_query(
        self, status: str | None, error_message: str | None, experiment_run_id: int
    ) -> tuple[str, dict[str, Any]]:
        """Build SQL query for status update with error message."""
        if not isinstance(status, str) or not status:
            raise ValueError('status is required for STATUS_WITH_ERROR update type')

        if not isinstance(error_message, str) or not error_message:
            raise ValueError('error_message is required for STATUS_WITH_ERROR update type')

        # Truncate error message if too long
        truncated_error = error_message[:1024] if len(error_message) > 1024 else error_message

        sql_query = """
            UPDATE experiment_run
            SET status = :status, error_message = :error_message, updated_at = :updated_at
            WHERE id = :experiment_run_id
        """

        query_params = {
            'status': status,
            'error_message': truncated_error,
            'updated_at': datetime.now(UTC),
            'experiment_run_id': experiment_run_id,
        }

        return sql_query, query_params

    def _build_model_saved_query(
        self, run_name: str | None, status: str | None, experiment_run_id: int
    ) -> tuple[str, dict[str, Any]]:
        """Build SQL query for model saved update."""
        if not isinstance(run_name, str) or not run_name:
            raise ValueError('run_name is required for MODEL_SAVED update type')

        if not isinstance(status, str) or not status:
            raise ValueError('status is required for MODEL_SAVED update type')

        sql_query = """
            UPDATE experiment_run
            SET run_name = :run_name, status = :status, updated_at = :updated_at
            WHERE id = :experiment_run_id
        """

        query_params = {
            'run_name': run_name,
            'status': status,
            'updated_at': datetime.now(UTC),
            'experiment_run_id': experiment_run_id,
        }

        return sql_query, query_params

    def _get_update_query_and_params(
        self, update_type: str, experiment_run_id: int, input_data: dict[str, Any]
    ) -> tuple[str, dict[str, Any]]:
        """Get SQL query and parameters based on update type."""
        status = input_data.get('status')
        error_message = input_data.get('error_message')
        run_name = input_data.get('run_name')

        if update_type == UpdateType.STATUS:
            return self._build_status_update_query(status, experiment_run_id)

        if update_type == UpdateType.STATUS_WITH_ERROR:
            return self._build_status_with_error_query(status, error_message, experiment_run_id)

        if update_type == UpdateType.MODEL_SAVED:
            return self._build_model_saved_query(run_name, status, experiment_run_id)

        raise ValueError(f'Invalid update_type: {update_type}')

    @activity.defn(name='update_experiment_run')
    def update_experiment_run(self, input_data: dict[str, Any]) -> None:
        """
        Update experiment run with status, errors, or model information.

        This activity provides a unified interface for all experiment run updates,
        supporting different update types through a single method. It automatically
        selects the appropriate SQL query based on the update type and parameters.

        Args:
            input_data: Configuration for experiment run update operation
                Required keys:
                    - metadata (dict): Workflow execution metadata
                    - experiment_run_id (int): Unique identifier for the experiment run
                    - update_type (str): Type of update (status, status_with_error, model_saved)
                Optional keys:
                    - status (str): New status for the experiment run
                    - error_message (str): Error message if update failed
                    - run_name (str): MLFlow run name if model was saved

        Raises:
            ValueError: If required parameters are missing for the update type
            RuntimeError: If update operation fails
        """
        metadata = input_data.get('metadata')
        experiment_run_id = input_data['experiment_run_id']
        update_type = input_data['update_type']
        status = input_data.get('status')

        try:
            sql_query, query_params = self._get_update_query_and_params(
                update_type, experiment_run_id, input_data
            )

            result = self._execute_update(sql_query, query_params)

            if result.get('rowcount', 0) == 0:
                error_msg = (
                    f'No experiment_run row updated for id={experiment_run_id} '
                    f'(row missing or id mismatch). update_type={update_type!r}, status={status!r}.'
                )
                raise ValueError(error_msg)

            self.info(
                f'Successfully updated experiment run {experiment_run_id} with status {status}',
                metadata,
            )
        except Exception as e:  # noqa: BLE001
            error_msg = f'Error updating experiment run - ID: {experiment_run_id}, Status: {status}, Error: {str(e)}'
            trace = traceback.format_exc()

            self.send_notification(
                metadata=metadata or {},
                notification_id='UPDATE_EXPERIMENT_RUN_ERROR',
                message=error_msg,
                block='update_experiment_run',
                level=NotificationLevel.ERROR,
                attachment_content=trace,
            )
            raise RuntimeError(error_msg) from e