from temporalio import activity, workflow

with workflow.unsafe.imports_passed_through():
    import tempfile
    import traceback
    from datetime import datetime
    from pathlib import Path
    from shutil import rmtree
    from typing import Any

    import mlflow
    import numpy as np
    import pandas as pd
    from pandas import DataFrame, to_datetime
    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_sync import MinioRepository
    from sientia_do.temporal.constants import (
        DATETIME_FORMAT,
        DATETIME_FORMAT_MS_WITH_TZ,
        DATETIME_FORMAT_WITH_TZ,
        now,
    )
    from sientia_do.utils.formatters import create_sample_dict
    from sientia_model.model_repository.mlflow_repository import SientiaMLflowRepository
    from sientia_model.model_repository.plugin_store import PluginStore

    from laborious.utils.dataframe_debug import build_dataframe_debug_message
    from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload


class MLFlow(SientiaMonitoring):
    """
    Temporal activities that talk to MLflow through ``SientiaMLflowRepository`` and ``SientiaModel`` wrappers.

    Models are resolved by registered name and the ``production`` alias (not by legacy stages or
    separate transform/predict flavors). ``get_cached_model`` loads or reuses a wrapper; inference
    uses ``wrapper.transform`` / ``wrapper.predict``; retrain uses ``wrapper.retrain`` or
    ``wrapper.train`` plus ``store_model`` and registry promotion via ``promote_to_alias``.

    Large inputs and outputs flow through ``MinioDataFramePayload`` when workflows offload parquet
    to MinIO. On failure, transform/predict still return a payload with ``success: False`` and
    error details for downstream gates.

    Attributes:
        mlflow_repository: Client for tracking, registry, artifact download, and run lifecycle.
        plugin_store: Reference to the store (runtime is installed on the worker; reserved for
            future store-backed helpers).
    """

    _MAX_DEBUG_DATAFRAME_ROWS = 100
    _DEFAULT_MODEL_ALIAS = 'production'

    def __init__(
        self,
        mlflow_repository: SientiaMLflowRepository,
        plugin_store: PluginStore,
        minio_repository: MinioRepository | None = None,
        logger: Logger | None = None,
        notification_handler: NotificationHandler | None = None,
        metrics_controller: MetricsController | None = None,
    ):
        """
        Attach shared MLflow and MinIO clients used by all ML activities in this mixin.

        Args:
         - mlflow_repository: Repository built by ``Activities`` (or injected in tests).
         - plugin_store: Plugin store instance from worker bootstrap.
         - minio_repository: MinIO client for ``MinioDataFramePayload`` upload/download.
         - logger: Structured logger.
         - notification_handler: Notifications on hard failures where applicable.
         - metrics_controller: Shared metrics controller.

        Return:
            None
        """

        self.minio_repository = minio_repository
        SientiaMonitoring.__init__(
            self,
            logger=logger,
            notification_handler=notification_handler,
            metrics_controller=metrics_controller,
        )
        self.mlflow_repository = mlflow_repository
        self.plugin_store = plugin_store

    def close(self) -> None:
        """
        Release MinIO manager resources held by the mixin.

        Return:
            None
        """
        if self.minio_repository is not None:
            try:
                self.minio_repository.close()
            finally:
                self.minio_repository = None
        SientiaMonitoring.shutdown(self)

    def __del__(self):
        self.close()

    def _debug_dataframe(self, message: str, data: Any, metadata: dict[str, Any]) -> None:
        """
        Log dataframe content only when row count is below the configured threshold

        Args:
         - message (str): Base log message to identify the dataframe in logs
         - data (Any): Dataframe-like object expected to expose shape and to_csv
         - metadata (dict[str, Any]): Workflow metadata for contextual logging
        """
        self.debug(
            build_dataframe_debug_message(
                message=message,
                data=data,
                max_rows=self._MAX_DEBUG_DATAFRAME_ROWS,
            ),
            metadata,
        )

    def _detect_and_parse_datetime_index(self, data: pd.DataFrame, metadata: dict) -> pd.DataFrame:
        """
        Ensure the transform output index is homogeneous and encoded as ``DATETIME_FORMAT_WITH_TZ`` strings.

        Accepts an all-string index (validated against the format), or all-``datetime`` /
        ``Timestamp`` (naive timestamps are localized to UTC before formatting). Mixed element types
        or unsupported types raise ``ValueError`` with a message logged at info level.

        Args:
         - data: DataFrame whose index carries the time dimension after transform.
         - metadata: Workflow metadata for log correlation.

        Return:
            ``pd.DataFrame``: Same frame with a normalized string index; empty frames are returned as-is.
        """

        if data.empty:
            self.info('Data is empty, skipping datetime index detection and parsing', metadata)
            return data

        index = data.index
        index_type = type(index[0])

        self.info(f'Index type: {index_type}', metadata)

        message = (
            f'Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, '
            f'string in format {DATETIME_FORMAT_WITH_TZ}.'
        )

        if not all(isinstance(i, index_type) for i in index):
            types = map(str, map(type, index))
            raise ValueError(f'{message}. Elements are {",".join(types)}')

        if index_type is str:
            try:
                pd.to_datetime(data.index, format=DATETIME_FORMAT_WITH_TZ)
            except ValueError as e:
                raise ValueError(f'{message}. Unable to parse given date format: {e}') from e

        elif index_type is datetime or index_type is pd.Timestamp:
            idx = data.index
            if hasattr(idx, 'tz') and idx.tz is None:
                data.index = idx.tz_localize('UTC')  # type: ignore[attr-defined]

            data.index = data.index.strftime(DATETIME_FORMAT_WITH_TZ)  # type: ignore[attr-defined]
        else:
            raise ValueError(f'{message}. Got {index_type}.')

        return data

    def _resolve_model_version_for_run(self, run_id: str) -> str:
        """
        Map an MLflow ``run_id`` to the latest registered model version that produced that run.

        ``search_model_versions`` may return multiple versions if the model was registered more than
        once for the same run; the highest numeric ``version`` wins so promotion targets the newest
        artifact set.

        Args:
         - run_id: Run UUID from ``retrain_model`` / experiment payload.

        Return:
            str: Registry version string acceptable by ``promote_to_alias``.

        Raises:
            ValueError: If the filter returns no versions (model not registered for this run).
        """

        versions = self.mlflow_repository._client.search_model_versions(
            filter_string=f"run_id='{run_id}'"
        )
        if not versions:
            raise ValueError(f'No registered model version found for run_id={run_id}')
        latest = max(versions, key=lambda v: int(v.version))
        return str(latest.version)

    def _resolve_model_alias(self, model_config: dict[str, Any] | None = None) -> str:
        """
        Resolve which MLflow alias should be used for model lookup/promotion.

        Args:
         - model_config: Optional model configuration that may include ``alias``.

        Return:
            str: Alias name trimmed and normalized; defaults to ``production``.
        """
        if not model_config:
            return self._DEFAULT_MODEL_ALIAS
        alias = str(model_config.get('alias', self._DEFAULT_MODEL_ALIAS)).strip()
        return alias or self._DEFAULT_MODEL_ALIAS

    @activity.defn(name='request_transform')
    def request_transform(self, input_data: dict[str, Any]) -> MinioDataFramePayload:
        """
        Pivot long-format sensor rows, load the production wrapper, and run ``wrapper.transform``.

        Expected tabular shape after load: columns including ``variable``, ``timestamp``, ``value``,
        and ``created_at`` for deduplication. Data are sorted by ``created_at``, de-duplicated per
        ``(variable, timestamp)``, pivoted wide, then passed to the model. ``model_config`` may
        include ``retention_minutes`` for wrapper cache TTL.

        Args:
         - input_data: Dict with ``metadata``, ``model_name``, ``data`` (``MinioDataFramePayload``
           dict or inline dataframe dict), and optional ``model_config``.

        Return:
            ``MinioDataFramePayload`` with transformed frame and ``success: True``, or a payload
            with ``success: False`` and exception details in ``status`` if transform fails.
        """
        metadata = input_data['metadata']
        self.info('Transforming data...', metadata)

        payload = MinioDataFramePayload.from_dict(input_data['data'])
        data = payload.retrieve(self.minio_repository, metadata)

        model_name = input_data['model_name']
        model_config = input_data.get('model_config', {})
        model_alias = self._resolve_model_alias(model_config)

        self._debug_dataframe('Raw input data:', data, metadata)

        # Long → wide: keep newest row per (variable, timestamp), then pivot for the wrapper API.
        data = data.sort_values('created_at', ascending=False).drop_duplicates(
            subset=['variable', 'timestamp'], keep='first'
        )

        data = data.pivot(index='timestamp', columns='variable', values='value')
        data.fillna(np.nan, inplace=True)

        data.columns.name = None
        data.index.name = None

        data['timestamp'] = data.index

        self._debug_dataframe('Processed input data:', data, metadata)

        try:
            wrapper = self.mlflow_repository.get_cached_model(
                model_name=model_name,
                alias=model_alias,
                retention_minutes=model_config.get('retention_minutes', 0),
                metadata=metadata,
            )
            transformed_df, transform_meta = wrapper.transform(data)
            if transform_meta:
                self.info(f'Wrapper transform metadata: {transform_meta}', metadata)

            transformed_df = self._detect_and_parse_datetime_index(transformed_df, metadata)

            response_data: dict[str, Any] = {'success': True, 'content': transformed_df}
        except Exception as e:
            response_data = {
                'success': False,
                'content': {'message': str(e), 'traceback': traceback.format_exc()},
            }

        self.debug(
            f'Transform raw response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
            metadata,
        )

        self.debug(
            f'Transform response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
            metadata,
        )

        self.info('Data transformed successfully', metadata)

        if not response_data.get('success', False):
            return MinioDataFramePayload.from_dataframe(
                dataframe=None,
                minio_repo=self.minio_repository,
                model_name=model_name,
                operation='transform',
                status=response_data,
                workflow_metadata=metadata,
                last_timestamp=payload.last_timestamp,
                logger=self.logger,
            )

        return MinioDataFramePayload.from_dataframe(
            dataframe=response_data['content'],
            minio_repo=self.minio_repository,
            model_name=model_name,
            operation='transform',
            workflow_metadata=metadata,
            status={
                'success': True,
            },
            last_timestamp=payload.last_timestamp,
            logger=self.logger,
        )

    @activity.defn(name='request_predict')
    def request_predict(self, input_data: dict[str, Any]) -> MinioDataFramePayload:
        """
        Load the production wrapper and call ``wrapper.predict`` on the prepared feature frame.

        The activity normalizes ``NaN`` to ``None`` for JSON-friendly columns, sets the row index
        the same way as ``retrain_model`` (UTC ``DatetimeIndex`` from ``DATETIME_FORMAT_WITH_TZ``),
        restores that index on the prediction frame, and records ``response_time``. Non-DataFrame
        predictions are coerced to a single ``prediction`` column.

        Args:
         - input_data: Same envelope as ``request_transform`` (``metadata``, ``model_name``,
           ``data``, optional ``model_config`` with ``retention_minutes``).

        Return:
            ``MinioDataFramePayload`` with predictions or error status mirroring transform behaviour.
        """
        metadata = input_data['metadata']
        self.info('Predicting data...', metadata)

        payload = MinioDataFramePayload.from_dict(input_data['data'])
        data = payload.retrieve(self.minio_repository, metadata)

        model_name = input_data['model_name']
        model_config = input_data.get('model_config', {})
        model_alias = self._resolve_model_alias(model_config)

        self._debug_dataframe('Input data for prediction:', data, metadata)

        data.replace(np.nan, None, inplace=True)

        data.index = pd.DatetimeIndex(
            to_datetime(data.index, format=DATETIME_FORMAT_WITH_TZ, utc=True)
        )
        input_index = data.index

        try:
            wrapper = self.mlflow_repository.get_cached_model(
                model_name=model_name,
                alias=model_alias,
                retention_minutes=model_config.get('retention_minutes', 0),
                metadata=metadata,
            )
            start_time = datetime.now()
            predict_data, pred_meta = wrapper.predict({}, data)
            end_time = datetime.now()

            if pred_meta:
                self.info(f'Wrapper predict metadata: {pred_meta}', metadata)

            if isinstance(predict_data, DataFrame):
                self._debug_dataframe(
                    'Data received from model prediction:', predict_data, metadata
                )
                predict_data.columns = pd.Index(['prediction'])
            else:
                self.debug(
                    f'Data received from model prediction (not a DataFrame): {predict_data}',
                    metadata,
                )
                predict_data = pd.DataFrame(predict_data, columns=['prediction'])

            predict_data.index = input_index
            predict_data['response_time'] = (end_time - start_time).total_seconds()

            response_data: dict[str, Any] = {'success': True, 'content': predict_data}
        except Exception as e:
            response_data = {
                'success': False,
                'content': {'message': str(e), 'traceback': traceback.format_exc()},
            }

        self.debug(
            f'Prediction response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
            metadata,
        )

        self.info('Data predicted successfully', metadata)

        if not response_data.get('success', False):
            return MinioDataFramePayload.from_dataframe(
                dataframe=None,
                minio_repo=self.minio_repository,
                model_name=model_name,
                operation='predict',
                status=response_data,
                workflow_metadata=metadata,
                last_timestamp=payload.last_timestamp,
                logger=self.logger,
            )

        return MinioDataFramePayload.from_dataframe(
            dataframe=response_data['content'],
            minio_repo=self.minio_repository,
            model_name=model_name,
            operation='predict',
            workflow_metadata=metadata,
            status={
                'success': True,
            },
            last_timestamp=payload.last_timestamp,
            logger=self.logger,
        )

    @activity.defn(name='retrain_model')
    def retrain_model(self, input_data: dict[str, Any]) -> dict[str, Any]:
        """
        Fit an updated wrapper from historical data, then log and register in MLflow.

        Flow: load long-format data from MinIO → dedupe/pivot like inference prep → require
        ``model_config['target']`` → read current ``production`` version for ``source_run_id`` tag →
        run ``wrapper.retrain`` outside run timing → ``start_run`` with retrain tags → log input
        CSV artifact → ``store_model`` and ``log_params``. Does not promote; the workflow calls
        ``update_production_model`` after validation.

        Args:
         - input_data: Must include ``metadata``, ``model_name``, ``data`` (payload), and
           ``model_config`` with at least ``target``.

        Return:
            On success: ``success``, ``experiment`` (``run_id``, ``experiment_id``, ``experiment_name``),
            ``message``, ``timestamp``. On failure: ``success: False``, error fields, and optional trace.
        """

        if self.minio_repository is None:
            raise ValueError('Minio repository not initialized')

        metadata = input_data['metadata']

        try:
            payload = MinioDataFramePayload.from_dict(input_data['data'])
            data = payload.retrieve(self.minio_repository, metadata)

        except Exception as e:
            trace = traceback.format_exc()
            self.send_notification(
                metadata=metadata,
                notification_id='ERROR_LOADING_RETRAIN_DATA',
                message=f'Error loading retrain data: {e}',
                block='retrain_model',
                level=NotificationLevel.ERROR,
                attachment_content=trace,
            )
            self.error(trace, metadata)
            return {
                'success': False,
                'message': f'Error loading retrain data: {e}',
                'traceback': trace,
                'timestamp': now().strftime(DATETIME_FORMAT_MS_WITH_TZ),
            }

        self.debug(f'Retrain data loaded successfully: shape {data.shape}', metadata)

        model_name = input_data['model_name']
        model_config = input_data.get('model_config', {})

        self.info(f'Retraining model {model_name}...', metadata)

        timestamp = data['timestamp'].max()
        self.debug(f'Timestamp: {timestamp}', metadata)

        if 'created_at' in data.columns:
            data = data.sort_values('created_at', ascending=False).drop_duplicates(
                subset=['variable', 'timestamp'], keep='first'
            )
        else:
            data = data.drop_duplicates(subset=['variable', 'timestamp'], keep='first')

        data.drop(columns=['model_id'], inplace=True, errors='ignore')
        data.drop(columns=['created_at'], inplace=True, errors='ignore')

        data = data.pivot(index='timestamp', columns='variable', values='value')
        data.fillna(np.nan, inplace=True)
        data.columns.name = None
        data.index.name = None

        data.index = pd.DatetimeIndex(
            to_datetime(data.index, format=DATETIME_FORMAT_WITH_TZ, utc=True)
        )

        target = model_config.get('target')
        if target is None:
            msg = 'model_config must include "target" for retraining'
            self.info(msg, metadata)
            return {
                'success': False,
                'experiment': None,
                'message': msg,
                'traceback': '',
                'timestamp': str(timestamp),
            }

        try:
            model_alias = self._resolve_model_alias(model_config)
            mv_src = self.mlflow_repository._client.get_model_version_by_alias(
                name=model_name,
                alias=model_alias,
            )
            source_run_id = mv_src.run_id

            wrapper = self.mlflow_repository.get_cached_model(
                model_name=model_name,
                alias=model_alias,
                retention_minutes=0,
                metadata=metadata,
            )

            # Keep heavy model fitting outside MLflow run timing.
            wrapper.retrain(data)

            run_name = f'{model_name}-retrain-{datetime.now().strftime("%Y%m%d%H%M%S")}'

            with self.mlflow_repository.start_run(
                model_name=model_name,
                run_name=run_name,
                experiment_name=model_name,
                tags={'retrain': 'true', 'source_run_id': source_run_id},
                metadata=metadata,
            ) as run_info:
                tmp_dir = tempfile.mkdtemp(prefix='laborious_retrain_')
                try:
                    raw_csv = Path(tmp_dir) / 'retrain_input.csv'
                    data.to_csv(raw_csv, index=False)
                    mlflow.log_artifact(str(raw_csv))
                finally:
                    rmtree(tmp_dir, ignore_errors=True)

                wrapper.store_model(name=model_name)

                self.mlflow_repository.log_params(
                    {
                        'retrain': 'true',
                        'retrain_date': datetime.now().isoformat(),
                        'source_run_id': source_run_id,
                        'retrain_samples': str(data.shape),
                    }
                )

                experiment_payload = {
                    'run_id': run_info.run_id,
                    'experiment_id': run_info.experiment_id,
                    'experiment_name': model_name,
                }

            return {
                'success': True,
                'experiment': experiment_payload,
                'message': 'Model retrained successfully.',
                'timestamp': str(timestamp),
            }

        except Exception as e:
            error_msg = f'Error retraining model {model_name}: {e}'
            self.info(error_msg, metadata)
            return {
                'success': False,
                'experiment': None,
                'message': error_msg,
                'traceback': traceback.format_exc(),
                'timestamp': str(timestamp),
            }

    @activity.defn(name='update_production_model')
    def update_production_model(self, input_data: dict[str, Any]) -> dict[Any, Any]:
        """
        Point the ``production`` alias at the model version registered for the retrain run.

        Resolves the highest numeric registry version whose ``run_id`` matches
        ``experiment['run_id']``, then calls ``promote_to_alias``. On failure, sends a notification
        and re-raises so the workflow can surface the error.

        Args:
         - input_data: ``metadata``, ``model_name``, and ``experiment`` with ``run_id`` and
           ``experiment_id`` (as returned from ``retrain_model``).

        Return:
            Dict with ``model_name``, promoted ``version``, ``mlflow_run_id``, ``mlflow_experiment_id``.
        """
        metadata = input_data['metadata']
        model_name = input_data['model_name']
        experiment = input_data['experiment']
        self.info(
            f'Updating production model {model_name} from experiment {experiment}...', metadata
        )

        try:
            run_id = experiment['run_id']
            experiment_id = experiment['experiment_id']

            version = self._resolve_model_version_for_run(run_id)

            promote_alias = self._resolve_model_alias(input_data.get('model_config'))
            self.mlflow_repository.promote_to_alias(
                model_name=model_name,
                version=version,
                alias=promote_alias,
                metadata=metadata,
            )

            self.info(f'Production model {model_name} updated successfully', metadata)
            return {
                'model_name': model_name,
                'version': version,
                'mlflow_run_id': run_id,
                'mlflow_experiment_id': experiment_id,
            }

        except Exception as e:
            trace = traceback.format_exc()
            self.send_notification(
                metadata=metadata,
                notification_id='UPDATE_PRODUCTION_MODEL_ERROR',
                message=f'Error updating production model {model_name}: {e}',
                block='update_production_model',
                level=NotificationLevel.ERROR,
                attachment_content=trace,
            )
            self.error(trace, metadata=metadata)
            raise e

    @activity.defn(name='get_reference_data')
    def get_reference_data(self, input_data: dict[str, Any]) -> list[dict] | None:
        """
        Download ``evaluation_data.csv`` from the MLflow run linked to ``production`` and parse it.

        Used by drift workflows to compare live data against the reference distribution logged with
        the model. Artifacts are downloaded to a temp directory, discovered via ``rglob`` (nested
        layout-safe), then timestamps are normalized to ``DATETIME_FORMAT`` string columns before
        returning record-oriented dicts.

        Args:
         - input_data: ``metadata`` and ``model_name`` for registry lookup.

        Return:
            List of row dicts, or ``None`` if the artifact path is missing or any step fails.
        """
        metadata = input_data['metadata']
        model_name = input_data['model_name']

        try:
            model_alias = self._resolve_model_alias(input_data.get('model_config'))
            mv = self.mlflow_repository._client.get_model_version_by_alias(
                name=model_name,
                alias=model_alias,
            )
            run_id = mv.run_id

            tmpdir = tempfile.mkdtemp(prefix='laborious_eval_')
            try:
                self.mlflow_repository.download_artifacts(
                    run_id=run_id,
                    artifact_path='evaluation_data.csv',
                    dst_path=tmpdir,
                    metadata=metadata,
                )
                csv_candidates = list(Path(tmpdir).rglob('evaluation_data.csv'))
                if not csv_candidates:
                    self.warning(f'Reference data not found for model {model_name}', metadata)
                    return None

                reference_data = pd.read_csv(csv_candidates[0])
            finally:
                rmtree(tmpdir, ignore_errors=True)

            reference_data['timestamp'] = to_datetime(reference_data['timestamp'])
            reference_data['timestamp'] = reference_data['timestamp'].dt.strftime(DATETIME_FORMAT)

            return reference_data.to_dict(orient='records')

        except Exception as e:
            self.warning(f'Reference data not found for model {model_name}: {e}', metadata)
            return None