from temporalio import activity, workflow

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

    from bson import ObjectId
    from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
    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.mongodb_repository_sync import MongoDBRepository
    from sientia_do.temporal.constants import now

    from laborious.activities.model_metrics import (
        next_adaptive_state,
        recovery_runs,
        threshold_breaches,
    )


class MongoActivities(SientiaMonitoring):
    """
    Activities that read from or write to MongoDB, sharing one `MongoDBRepository`.

    - `write_prediction_metadata`: centralized, best-effort writer for prediction gate metadata
      (triggered filter, policy, counter). Gates only compute and return a metadata document;
      they never write it themselves. Keeping the write here, instead of in each gate, puts the
      one real write point in the subworkflow that's about to persist a `predictions` row — the
      only place that reliably knows whether one is being written at all — and lets other
      metadata sources reuse the same activity instead of duplicating write logic wherever they
      originate.
    - `update_simple_metrics_adaptive_state`: adaptive frequency state on the `pipelines` doc of
      a simple metrics schedule (SIENTIAPDE-1995).
    """

    mongo_repository: MongoDBRepository | None = None
    _PREDICTION_METADATA_COLLECTION = 'prediction_metadata'
    _PIPELINES_COLLECTION = 'pipelines'

    def __init__(
        self,
        mongo_repository: MongoDBRepository | None = None,
        logger: Logger | None = None,
        notification_handler: NotificationHandler | None = None,
        metrics_controller: MetricsController | None = None,
    ):
        """
        Initialize the Mongo activities.

        Args:
            mongo_repository: MongoDB, where gate metadata (triggered filter, policy, counter) is
                written best-effort and the simple metrics adaptive state is kept. `None` is a
                valid state: both activities become no-ops.

        Raises:
            Exception: If BaseActivity initialization fails
        """
        self.mongo_repository = mongo_repository
        SientiaMonitoring.__init__(
            self,
            logger=logger,
            notification_handler=notification_handler,
            metrics_controller=metrics_controller,
        )

    def close(self) -> None:
        """
        Close the Mongo repository and clean up resources.
        """
        if self.mongo_repository is not None:
            try:
                self.mongo_repository.close()
            finally:
                self.mongo_repository = None
        SientiaMonitoring.shutdown(self)

    def __del__(self):
        self.close()

    @activity.defn(name='write_prediction_metadata')
    def write_prediction_metadata(self, input_data: dict[str, Any]) -> str | None:
        """
        Best-effort write of one or more gate metadata documents to Mongo, combined into a single
        `prediction_metadata` document.

        Implements both business rules for this feature in one place:
        - The prediction write must never be blocked or failed by a metadata write. A failed
          `insert(...)` is swallowed, never re-raised, logged here with call-site context (which
          `mongo_repository.insert(...)` itself can't know) — matching how `storage.py`'s
          `cleanup_minio_objects_expired` also logs again at the call site even though the
          repository already logs internally and emits its own Prometheus metric
          (`MONGODB_WRITE_ERROR_COUNT`); the generic library log is for the metric/trace, this
          one is for whoever is reading logs for this specific call.
        - Nothing is written when there's nothing relevant to report. An empty/missing
          `documents` list (or one full of falsy/empty entries) skips Mongo entirely, so callers
          can call this unconditionally and get `None` back for "no metadata this time" instead
          of deciding upfront whether there's something worth saving.

        Takes a *list* of documents, not a single one: each caller contributes its own document
        (shaped like `{'filters': [...]}`), and this flattens every `filters` entry across all of
        them into one `prediction_metadata` document before the single insert. Today only one
        gate ever contributes a document per prediction cycle (the gate chain short-circuits on
        the first non-`None` `path_flag`), but the list shape means this activity won't need to
        change again once more than one source contributes metadata for the same cycle.

        `insert(...)` doesn't hand back the inserted document's `_id`, so it is generated here,
        client-side, before the insert, as a plain string (not a BSON ObjectId) — matching this
        codebase's convention of string `_id`s (e.g. the `models` collection) and the 24-hex-char
        size of the `predictions.metadata_id` Postgres column that will reference it.

        Args:
            input_data: Dict with:
                - documents (list[dict[str, Any]]): Metadata payloads to combine, each shaped
                  like `{'filters': [...]}`. Missing/empty, or entries with no `filters`, skip
                  the write entirely.
                - metadata (dict): Workflow metadata for contextual logging.

        Return:
            The generated id, or `None` when there's nothing to write, Mongo isn't configured, or
            the insert failed.
        """
        metadata = input_data['metadata']
        documents = input_data.get('documents') or []

        filters = [
            entry for document in documents if document for entry in document.get('filters', [])
        ]

        if not filters or self.mongo_repository is None:
            return None

        generated_id = str(ObjectId())

        try:
            self.mongo_repository.insert(
                self._PREDICTION_METADATA_COLLECTION,
                {'filters': filters, '_id': generated_id},
                metadata,
            )
        except Exception as e:
            self.error(f'Failed to write prediction metadata: {e}', metadata)
            return None

        return generated_id

    @activity.defn(name='update_simple_metrics_adaptive_state')
    def update_simple_metrics_adaptive_state(self, input_data: dict[str, Any]) -> dict | None:
        """
        Update the adaptive frequency state on the schedule's `pipelines` doc.

        The orchestrator derives the schedule frequency from `adaptive_state.escalated` and
        `frequency_alert`. State is read from Mongo, not from the workflow input, because the
        input only refreshes on the orchestrator tick and would be stale while escalated.
        `updated_at` is bumped only when `escalated` flips: that is what makes the orchestrator
        rewrite the schedule, and a streak change alone does not change the frequency.

        Args:
            input_data: Dict with:
                - metadata (dict): Workflow metadata; `schedule_name` identifies the doc.
                - simple_metrics (list[dict]): Rows from `calculate_simple_metrics`; empty for an
                  empty or under-sampled window, which counts as a non-breach run.
                - thresholds (dict): `simple_metrics_thresholds` from `model_config` (may be None: no breach).
                - interval_minutes (int): Window size, used to compute the recovery runs.

        Return:
            The new state, or `None` when adaptive mode is off, Mongo is not configured, the doc
            is missing or `frequency_alert` is invalid.

        Raises:
            Exception: Mongo read/write failures, so the workflow run fails explicitly.
        """
        metadata = input_data['metadata']
        if self.mongo_repository is None:
            return None

        # Mongo errors propagate on purpose: the run must fail explicitly so the problem shows
        # up in the Temporal metrics (PR #59 review). Metrics were already exported by then.
        query = {'schedule_name': metadata['schedule_name']}
        pipeline = self.mongo_repository.find_one(self._PIPELINES_COLLECTION, query, metadata)
        if not pipeline:
            # Likely misconfig: MONGODB_DATABASE_NAME differs from the orchestrator's, or the
            # schedule_name does not match its doc.
            self.warning(
                f'No pipelines doc for schedule {metadata["schedule_name"]}: '
                'adaptive state skipped',
                metadata,
            )
            return None
        frequency_alert = pipeline.get('frequency_alert')
        if not frequency_alert:
            return None

        try:
            runs = recovery_runs(input_data['interval_minutes'], str(frequency_alert))
        except (TypeError, ValueError):
            # Config error, not a failure: the orchestrator ignores an invalid `frequency_alert`
            # too (keeps `frequency`) and already sends SIMPLE_METRICS_FREQUENCY_ALERT_BELOW_FLOOR
            # for it. `interval_minutes` is untyped Mongo config as well (e.g. "60" or null).
            self.warning(
                f'Invalid frequency_alert={frequency_alert!r} or '
                f'interval_minutes={input_data["interval_minutes"]!r}: adaptive state skipped',
                metadata,
            )
            return None

        breached = bool(
            threshold_breaches(input_data['simple_metrics'], input_data.get('thresholds'))
        )
        state = pipeline.get('adaptive_state')
        if not isinstance(state, dict):
            # Missing or hand-edited garbage: not escalated, same as the orchestrator reads it.
            state = {'escalated': False, 'healthy_streak': 0}
        new_state = next_adaptive_state(state, breached, runs)
        if new_state == state:
            return new_state

        update: dict[str, Any] = {'adaptive_state': new_state}
        if new_state['escalated'] != bool(state.get('escalated')):
            # BSON datetime, never a string: the orchestrator normalizes `updated_at` with
            # `.replace(tzinfo=UTC).strftime(...)` and a string would fail its whole tick.
            update['updated_at'] = now()
        self.mongo_repository.update_one(
            self._PIPELINES_COLLECTION, query, {'$set': update}, metadata
        )

        self.info(f'Adaptive state of {query["schedule_name"]}: {state} -> {new_state}', metadata)
        return new_state