"""
End-to-end tests for the Drift workflow.

The drift suite drives the **real** ``sientia_model.analytics.drift_analysis.DriftAnalysis``
analyzer (no mocking). Each scenario exercises the full pipeline:

    laborious_data (Postgres)
        -> load_custom_query
        -> calculate_drift (DriftAnalysis univariate + multivariate)
        -> export_data_to_postgres (sientia_data.drift_metrics)

Coverage focus:

- Happy path persists every column required by ``sientia_data.drift_metrics``
  with a valid reference dataset downloaded from MLflow.
- 30% fallback path activates when the MLflow reference is unavailable and
  emits the ``MODEL_METRICS_REFERENCE_DATA_WARNING`` notification.
- Empty target data short-circuits the workflow without persisting anything.
- Invalid ``chunk_period`` is rejected by ``calculate_drift``.
- ``chunk_period='s'`` preserves second-level precision in
  ``chunk_start_date``.

Tests assert behavioral / structural properties (column presence, NOT NULL
constraints, business-key invariants) rather than exact numeric values, since
those depend on the real analyzer implementation and synthetic data.
"""

from datetime import datetime, timedelta, timezone
from pathlib import Path
from unittest.mock import MagicMock, patch

import pandas as pd
import pytest
from sqlalchemy import text
from temporalio.testing import WorkflowEnvironment
from temporalio.worker import Worker

from e2e.helpers import (
    insert_target_data_for_drift,
    load_scenario_input,
    make_workflow_id,
    start_and_await_workflow,
)
from laborious.activities.activities import Activities
from laborious.activities.model_metrics import ModelMetrics
from laborious.workflows.drift import Drift
from sientia_model.analytics.drift_analysis import DriftAnalysis

# Drift columns persisted on every row in ``sientia_data.drift_metrics`` —
# mirrors the production DDL.
EXPECTED_DRIFT_COLUMNS = [
    'id',
    'model_id',
    'feature',
    'method',
    'value',
    'alert',
    'chunk_index',
    'chunk_start_date',
    'chunk_end_date',
    'accurate',
    'timestamp',
    'created_at',
]

# Columns the DDL marks as NOT NULL. ``feature`` and ``timestamp`` are
# nullable in the production schema (multivariate rows do not bind to a
# single feature; ``timestamp`` is allowed to be empty when upstream data has
# no usable instant).
NON_NULL_DRIFT_COLUMNS = {
    'id',
    'model_id',
    'method',
    'value',
    'alert',
    'chunk_index',
    'chunk_start_date',
    'chunk_end_date',
    'accurate',
    'created_at',
}

DEFAULT_DRIFT_METHODS = ['kolmogorov_smirnov', 'jensen_shannon', 'wasserstein']


def _chunk_dataframe_skip_empty_groups(
    self: DriftAnalysis,
    df: pd.DataFrame,
    timestamp_col: str,
    chunk_period: str,
) -> list[tuple[int, pd.DataFrame]]:
    """
    Same as ``DriftAnalysis._chunk_dataframe`` but omit empty time buckets.

    ``pd.Grouper(freq='s')`` yields every second between min and max timestamp;
    empty buckets still appear in the groupby iterator and produce invalid
    drift rows (e.g. NaT timestamps) that ``calculate_drift`` later filters out
    entirely. Production fix belongs in ``sientia_model``; this shim keeps the
    e2e honest about second-level chunk boundaries with sparse samples.
    """
    grouped = df.groupby(pd.Grouper(key=timestamp_col, freq=chunk_period), dropna=True)
    chunks: list[tuple[int, pd.DataFrame]] = []
    idx = 0
    for _, chunk in grouped:
        if chunk.empty:
            continue
        chunks.append((idx, chunk.copy()))
        idx += 1
    return chunks


def _drift_input(model_id: int, **overrides) -> dict:
    """Load the base drift scenario JSON and apply ad-hoc overrides."""
    input_data = load_scenario_input('drift_base.json', model_id=model_id)
    input_data.update(overrides)
    return input_data


def _recent_minute_timestamps(count: int, offset_minutes: int = 6) -> list[str]:
    """
    Build ``count`` consecutive UTC minute timestamps placed in the recent past.

    The Drift workflow filters target rows with ``timestamp > NOW() - INTERVAL``,
    so timestamps must be recent for tests to retrieve any data. Snapping to
    minute precision keeps the helper deterministic regardless of clock skew.

    Args:
     - count (int): How many consecutive minute timestamps to generate.
     - offset_minutes (int): Minutes ago for the EARLIEST generated timestamp.

    Return:
     list[str]: ISO strings with ``+0000`` offset, one per minute.
    """
    base = datetime.now(timezone.utc).replace(second=0, microsecond=0) - timedelta(
        minutes=offset_minutes
    )
    return [
        (base + timedelta(minutes=i)).strftime('%Y-%m-%d %H:%M:%S%z') for i in range(count)
    ]


def _configure_reference_csv(mlflow_repository_stub, reference_rows: pd.DataFrame) -> None:
    """
    Wire ``mlflow_repository_stub`` so ``get_reference_data`` returns
    ``reference_rows`` by writing them to ``dst_path/evaluation_data.csv``.

    Args:
     - mlflow_repository_stub: External MLflow repository fixture.
     - reference_rows (pd.DataFrame): Rows to expose as the production reference.
    """

    def _download(run_id: str, artifact_path: str, dst_path: str, metadata=None):
        target = Path(dst_path) / 'evaluation_data.csv'
        reference_rows.to_csv(target, index=False)

    mlflow_repository_stub._client.get_model_version_by_alias.return_value = MagicMock(
        run_id='fake-reference-run'
    )
    mlflow_repository_stub.download_artifacts.side_effect = _download


def _force_reference_unavailable(mlflow_repository_stub) -> None:
    """Make ``get_reference_data`` return ``None`` by failing alias resolution."""
    mlflow_repository_stub._client.get_model_version_by_alias.side_effect = Exception(
        'no production alias registered'
    )


def _select_drift_rows(postgres_engine, model_id: int) -> list[dict]:
    """Read every persisted drift row for ``model_id`` ordered by chunk/feature/method."""
    with postgres_engine.connect() as conn:
        rows = (
            conn.execute(
                text(
                    'SELECT * FROM sientia_data.drift_metrics '
                    'WHERE model_id = :m '
                    'ORDER BY chunk_index, feature, method'
                ),
                {'m': str(model_id)},
            )
            .mappings()
            .all()
        )
    return [dict(row) for row in rows]


def _assert_required_columns_populated(rows: list[dict]) -> None:
    """Validate column presence and NOT NULL constraints on every row."""
    assert rows, 'expected at least one drift row to be persisted'
    seen_columns = set(rows[0].keys())
    for column in EXPECTED_DRIFT_COLUMNS:
        assert column in seen_columns, f'Missing drift column in postgres: {column}'
    for row in rows:
        for column in NON_NULL_DRIFT_COLUMNS:
            assert row[column] is not None, f"Column '{column}' is NULL in {row}"
        assert 'p_value' not in row, 'p_value must not be persisted to drift_metrics'


@pytest.mark.asyncio
@pytest.mark.integration
async def test_drift_happy_path_persists_all_columns_with_reference_data(
    temporal_test_env: WorkflowEnvironment,
    temporal_worker_drift: Worker,
    test_activities: Activities,
    postgres_engine,
    mlflow_repository_stub,
):
    """
    Scenario D.1.1: Happy path with reference data downloaded from MLflow.

    Drives the full pipeline against the real ``DriftAnalysis``. Asserts:

    - One row is persisted per ``(chunk_index, feature, method)`` combination
      plus the multivariate row block, with every column required by
      ``sientia_data.drift_metrics`` populated.
    - The three default univariate methods are forwarded to the analyzer.
    - ``model_id`` and ``timestamp`` are stamped by the activity (not by the
      analyzer); ``timestamp`` equals ``max(target_data.timestamp)`` and is
      identical on every persisted row.
    - ``chunk_start_date`` / ``chunk_end_date`` are persisted as ISO text so
      the analyzer's nanosecond-precision boundaries survive the ``text``
      column type.
    - ``accurate=True`` because the reference dataset was available.
    """
    client = temporal_test_env.client
    model_id = 411

    target_timestamps = _recent_minute_timestamps(count=10)
    insert_target_data_for_drift(
        postgres_engine,
        model_id=model_id,
        timestamps=target_timestamps,
        variables_values={
            'sensor_1': [10.0 + i * 0.1 for i in range(10)],
            'sensor_2': [20.0 + i * 0.5 for i in range(10)],
        },
    )

    reference_df = pd.DataFrame(
        {
            'timestamp': [
                f'2023-12-31 11:{minute:02d}:00+00:00' for minute in range(10)
            ],
            'sensor_1': [9.0 + i * 0.05 for i in range(10)],
            'sensor_2': [18.0 + i * 0.25 for i in range(10)],
        }
    )
    _configure_reference_csv(mlflow_repository_stub, reference_df)

    input_data = _drift_input(model_id)
    await start_and_await_workflow(
        client, Drift.run, input_data, make_workflow_id('test-drift-happy-path')
    )

    rows = _select_drift_rows(postgres_engine, model_id)

    exported_csv_path = '/tmp/test_drift_happy_path_exported.csv'
    pd.DataFrame(rows).to_csv(exported_csv_path, index=False)
    print(
        f'\n[test_drift_happy_path] Exported drift dataframe '
        f'({len(rows)} rows) -> {exported_csv_path}'
    )

    _assert_required_columns_populated(rows)

    # The activity drops the target column from the feature list, so only
    # ``sensor_2`` participates in univariate analysis (``sensor_1`` is the
    # configured target). Multivariate produces one row per chunk regardless.
    univariate_rows = [r for r in rows if r['feature'] != 'multivariate']
    multivariate_rows = [r for r in rows if r['feature'] == 'multivariate']
    assert univariate_rows, 'expected univariate drift rows for non-target features'
    assert multivariate_rows, 'expected one multivariate drift row per chunk'

    # All three default methods must reach the analyzer.
    assert {r['method'] for r in univariate_rows} == set(DEFAULT_DRIFT_METHODS)
    assert all(r['method'] == 'multivariate' for r in multivariate_rows)
    assert {r['feature'] for r in univariate_rows} == {'sensor_2'}

    # ``timestamp`` is stamped uniformly with ``max(target_data.timestamp)``.
    expected_timestamp = pd.to_datetime(max(target_timestamps), utc=True)
    persisted_timestamps = {pd.to_datetime(r['timestamp'], utc=True) for r in rows}
    assert len(persisted_timestamps) == 1, (
        'timestamp must be uniform across all drift rows '
        f'(got {len(persisted_timestamps)} distinct values)'
    )
    assert pd.Timestamp(persisted_timestamps.pop()) == expected_timestamp, (
        'timestamp must equal max(target_data.timestamp)'
    )

    # ``model_id`` is stamped by ``calculate_drift`` (not produced by the analyzer).
    assert all(r['model_id'] == str(model_id) for r in rows), (
        'model_id must be stamped on every drift row'
    )

    # Reference path → accurate=True.
    assert all(r['accurate'] is True for r in rows), (
        'reference path should mark all rows as accurate'
    )

    # ISO text serialization preserves ordering between start/end of each chunk.
    for row in rows:
        assert 'T' in row['chunk_start_date'], (
            f"chunk_start_date should be ISO text, got {row['chunk_start_date']!r}"
        )
        assert 'T' in row['chunk_end_date'], (
            f"chunk_end_date should be ISO text, got {row['chunk_end_date']!r}"
        )
        assert row['chunk_start_date'] <= row['chunk_end_date'], (
            f'chunk_start_date must precede chunk_end_date '
            f"(start={row['chunk_start_date']}, end={row['chunk_end_date']})"
        )


@pytest.mark.asyncio
@pytest.mark.integration
async def test_drift_uses_30pct_fallback_when_reference_unavailable(
    temporal_test_env: WorkflowEnvironment,
    temporal_worker_drift: Worker,
    test_activities: Activities,
    postgres_engine,
    mlflow_repository_stub,
    notification_inserts,
):
    """
    Scenario D.1.2: ``get_reference_data`` returns ``None`` (production alias
    missing), so ``calculate_drift`` falls back to using the first 30% of
    target rows as reference. Persisted rows must report ``accurate=False``
    and a ``MODEL_METRICS_REFERENCE_DATA_WARNING`` notification must be
    emitted to mongo.
    """
    client = temporal_test_env.client
    model_id = 412

    target_timestamps = _recent_minute_timestamps(count=10)
    insert_target_data_for_drift(
        postgres_engine,
        model_id=model_id,
        timestamps=target_timestamps,
        variables_values={
            'sensor_1': [10.0 + i * 0.1 for i in range(10)],
            'sensor_2': [20.0 + i * 0.5 for i in range(10)],
        },
    )

    _force_reference_unavailable(mlflow_repository_stub)

    input_data = _drift_input(model_id)
    await start_and_await_workflow(
        client, Drift.run, input_data, make_workflow_id('test-drift-fallback')
    )

    rows = _select_drift_rows(postgres_engine, model_id)
    _assert_required_columns_populated(rows)

    assert all(r['accurate'] is False for r in rows), (
        'fallback path must mark all rows as inaccurate'
    )

    fallback_warnings = [
        call
        for call in notification_inserts.call_args_list
        if call.args
        and isinstance(call.args[0], dict)
        and call.args[0].get('notification_id') == 'MODEL_METRICS_REFERENCE_DATA_WARNING'
    ]
    assert len(fallback_warnings) >= 1, 'expected reference fallback warning notification'


@pytest.mark.asyncio
@pytest.mark.integration
async def test_drift_empty_target_data_short_circuits_workflow(
    temporal_test_env: WorkflowEnvironment,
    temporal_worker_drift: Worker,
    test_activities: Activities,
    postgres_engine,
    mlflow_repository_stub,
):
    """
    Scenario D.3.1: When ``load_custom_query`` returns no rows the workflow
    must return early without invoking the analyzer or writing any drift rows.
    """
    client = temporal_test_env.client
    model_id = 431

    with postgres_engine.begin() as conn:
        conn.execute(text(f'DELETE FROM sientia_data.laborious_data WHERE model_id = {model_id}'))

    _force_reference_unavailable(mlflow_repository_stub)

    input_data = _drift_input(model_id)
    await start_and_await_workflow(
        client, Drift.run, input_data, make_workflow_id('test-drift-empty-target')
    )

    with postgres_engine.connect() as conn:
        count = conn.execute(
            text('SELECT COUNT(*) FROM sientia_data.drift_metrics WHERE model_id = :m'),
            {'m': str(model_id)},
        ).scalar()
    assert count == 0


@pytest.mark.asyncio
@pytest.mark.integration
async def test_drift_invalid_chunk_period_raises_value_error(
    temporal_test_env: WorkflowEnvironment,
    temporal_worker_drift: Worker,
    test_activities: Activities,
    postgres_engine,
    mlflow_repository_stub,
):
    """
    Scenario D.4.2: ``calculate_drift`` validates ``chunk_period`` and rejects
    anything other than ``min`` / ``s``. The workflow must surface the
    ``ValueError`` and persist nothing.
    """
    client = temporal_test_env.client
    model_id = 442

    target_timestamps = _recent_minute_timestamps(count=5)
    insert_target_data_for_drift(
        postgres_engine,
        model_id=model_id,
        timestamps=target_timestamps,
        variables_values={
            'sensor_1': [1.0, 2.0, 3.0, 4.0, 5.0],
            'sensor_2': [10.0, 20.0, 30.0, 40.0, 50.0],
        },
    )
    _force_reference_unavailable(mlflow_repository_stub)

    input_data = _drift_input(model_id, chunk_period='hour')

    with pytest.raises(Exception) as excinfo:
        await start_and_await_workflow(
            client,
            Drift.run,
            input_data,
            make_workflow_id('test-drift-bad-chunk-period'),
        )

    # Temporal wraps the activity ValueError in WorkflowFailureError; the
    # message may live on ``.message`` or ``str(exc)`` depending on the SDK
    # error class, so walk the cause chain looking for the guard text.
    cause_descriptions = []
    current: BaseException | None = excinfo.value
    while current is not None:
        cause_descriptions.append(
            getattr(current, 'message', None) or str(current) or repr(current)
        )
        current = current.__cause__
    assert any('Invalid chunk period' in msg for msg in cause_descriptions), (
        f'Expected ValueError about chunk period in chain, got: {cause_descriptions}'
    )

    with postgres_engine.connect() as conn:
        count = conn.execute(
            text('SELECT COUNT(*) FROM sientia_data.drift_metrics WHERE model_id = :m'),
            {'m': str(model_id)},
        ).scalar()
    assert count == 0


@pytest.mark.asyncio
@pytest.mark.integration
async def test_drift_empty_merge_skips_export_without_insufficient_notification(
    temporal_test_env: WorkflowEnvironment,
    temporal_worker_drift: Worker,
    test_activities: Activities,
    postgres_engine,
    mlflow_repository_stub,
):
    """
    Scenario D.4.3a: When the analyzer returns an empty merged frame (no metric rows),
    ``calculate_drift`` yields ``[]``; the workflow skips export. Real insufficient-data
    cases are signaled by ``DriftInsufficientDataError`` inside ``sientia_model``, not by
    empty output alone.
    """
    client = temporal_test_env.client
    model_id = 444

    target_timestamps = _recent_minute_timestamps(count=5)
    insert_target_data_for_drift(
        postgres_engine,
        model_id=model_id,
        timestamps=target_timestamps,
        variables_values={
            'sensor_1': [1.0 + i * 0.1 for i in range(5)],
            'sensor_2': [10.0 + i * 0.5 for i in range(5)],
        },
    )
    _force_reference_unavailable(mlflow_repository_stub)

    input_data = _drift_input(model_id, chunk_period='min')
    empty_merge = pd.DataFrame(
        columns=[
            'timestamp',
            'feature',
            'method',
            'value',
            'alert',
            'chunk_index',
            'chunk_start_date',
            'chunk_end_date',
            'threshold',
            'drift_type',
        ]
    )
    with patch.object(ModelMetrics, 'get_drift_metrics', return_value=empty_merge):
        await start_and_await_workflow(
            client,
            Drift.run,
            input_data,
            make_workflow_id('test-drift-empty-merge'),
        )

    with postgres_engine.connect() as conn:
        count = conn.execute(
            text('SELECT COUNT(*) FROM sientia_data.drift_metrics WHERE model_id = :m'),
            {'m': str(model_id)},
        ).scalar()
    assert count == 0


@pytest.mark.asyncio
@pytest.mark.integration
async def test_drift_chunk_period_seconds_sufficient_data_preserves_seconds_in_chunk_start_date(
    temporal_test_env: WorkflowEnvironment,
    temporal_worker_drift: Worker,
    test_activities: Activities,
    postgres_engine,
    mlflow_repository_stub,
):
    """
    Scenario D.4.3b: With enough sub-minute samples and ``chunk_period='s'``, drift rows
    persist and ``chunk_start_date`` keeps second-level precision (incl. second=30).

    ``DriftAnalysis._chunk_dataframe`` is patched to skip empty ``pd.Grouper(freq='s')``
    buckets so sparse seconds between samples do not flood the pipeline with NaT rows;
    the durable fix belongs in ``sientia_model``.
    """
    client = temporal_test_env.client
    model_id = 443

    base = datetime.now(timezone.utc).replace(second=0, microsecond=0) - timedelta(minutes=6)
    target_timestamps = []
    sensor_1_vals = []
    sensor_2_vals = []
    for minute_offset in range(6):
        t0 = base + timedelta(minutes=minute_offset)
        t1 = t0 + timedelta(seconds=30)
        target_timestamps.append(t0.strftime('%Y-%m-%d %H:%M:%S%z'))
        target_timestamps.append(t1.strftime('%Y-%m-%d %H:%M:%S%z'))
        v0 = 1.0 + minute_offset * 0.1
        v1 = v0 + 0.05
        sensor_1_vals.extend([v0, v1])
        sensor_2_vals.extend([10.0 + v0, 10.0 + v1])

    insert_target_data_for_drift(
        postgres_engine,
        model_id=model_id,
        timestamps=target_timestamps,
        variables_values={
            'sensor_1': sensor_1_vals,
            'sensor_2': sensor_2_vals,
        },
    )
    _force_reference_unavailable(mlflow_repository_stub)

    input_data = _drift_input(model_id, chunk_period='s')
    with patch.object(DriftAnalysis, '_chunk_dataframe', _chunk_dataframe_skip_empty_groups):
        await start_and_await_workflow(
            client,
            Drift.run,
            input_data,
            make_workflow_id('test-drift-chunk-seconds-sufficient'),
        )

    with postgres_engine.connect() as conn:
        rows = (
            conn.execute(
                text(
                    'SELECT chunk_start_date FROM sientia_data.drift_metrics '
                    'WHERE model_id = :m ORDER BY chunk_index'
                ),
                {'m': str(model_id)},
            )
            .mappings()
            .all()
        )

    assert rows, 'expected at least one drift row to be persisted'
    seconds_present = {pd.Timestamp(r['chunk_start_date']).second for r in rows}
    assert 30 in seconds_present, (
        f'expected at least one chunk_start_date with seconds=30, got {seconds_present}'
    )