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

The MLflow registry is fully stubbed because no real artifacts exist in a
test container; we only validate that the workflow:

- Loads training data via ``load_query_with_minio_offload``.
- Calls ``retrain_model`` with a payload pointing at MinIO.
- Calls ``update_production_model`` only when retrain succeeds.
- Persists ``sientia_data.log_retrain`` rows with all required columns;
  success rows carry the new ``version`` / ``mlflow_run_id`` /
  ``mlflow_experiment_id`` while failure rows leave them ``NULL``.

The production DDL drops the legacy ``id`` / ``created_at`` columns and
moves ``mlflow_experiment_id`` to ``int8`` and ``model_id`` to ``text``.
The stubs used here therefore emit ``experiment_id`` as an integer to fit
the new column type.
"""

from contextlib import contextmanager
from datetime import datetime, timedelta, timezone
from unittest.mock import MagicMock, patch

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.workflows.minimal_retrain import MinimalRetrain

# Columns defined by the production DDL for ``sientia_data.log_retrain``.
# The legacy ``retrain_reports`` table had ``id`` and ``created_at``; the new
# DDL drops both. ``mlflow_experiment_id`` is ``int8`` and ``model_id`` is
# ``text``.
EXPECTED_RETRAIN_REPORT_COLUMNS = [
    'mlflow_experiment_id',
    'mlflow_run_id',
    'model_id',
    'model_name',
    'status',
    'timestamp',
    'version',
]


def _retrain_input(model_id: int, **overrides) -> dict:
    """Load and override the minimal-retrain base scenario."""
    payload = load_scenario_input('minimal_retrain_base.json', model_id=model_id)
    payload.update(overrides)
    return payload


def _seed_retrain_training_rows(postgres_engine, model_id: int) -> None:
    """
    Insert training rows in long format that pivot cleanly into
    ``index=timestamp`` / ``columns={sensor_1, sensor_2}`` for ``retrain_model``.
    """
    target_timestamps = [
        (datetime.now(timezone.utc).replace(second=0, microsecond=0) - timedelta(minutes=10 - i))
        .strftime('%Y-%m-%d %H:%M:%S%z')
        for i in range(5)
    ]
    insert_target_data_for_drift(
        postgres_engine,
        model_id=model_id,
        timestamps=target_timestamps,
        variables_values={
            'sensor_1': [10.0, 11.0, 12.0, 13.0, 14.0],
            'sensor_2': [20.0, 21.0, 22.0, 23.0, 24.0],
        },
    )


def _configure_retrain_happy_path(mlflow_repository_stub) -> None:
    """
    Wire ``mlflow_repository_stub`` so retrain + update_production succeed.

    Mocks (in order of consumption):

    - ``_client.get_model_version_by_alias``: returns ``mv`` with a stable
      ``run_id`` (used as ``source_run_id``).
    - ``get_cached_model``: returns a wrapper exposing inert ``retrain`` and
      ``store_model`` methods.
    - ``start_run``: returns a context manager yielding a ``run_info`` with
      run/experiment ids.
    - ``log_params``: inert.
    - ``_client.search_model_versions``: returns one registry entry whose
      ``version`` is promoted by ``update_production_model``.
    - ``promote_to_alias``: inert success.
    """
    mv_src = MagicMock()
    mv_src.run_id = 'source-run-id'

    new_version = MagicMock()
    new_version.version = '7'
    new_version.run_id = 'retrain-run-id'

    mlflow_repository_stub._client.get_model_version_by_alias.return_value = mv_src

    cached_wrapper = MagicMock()
    cached_wrapper.retrain = MagicMock(return_value=None)
    cached_wrapper.store_model = MagicMock(return_value=None)
    mlflow_repository_stub.get_cached_model.return_value = cached_wrapper

    @contextmanager
    def fake_start_run(**kwargs):
        run_info = MagicMock()
        run_info.run_id = 'retrain-run-id'
        # ``mlflow_experiment_id`` is ``int8`` in the new DDL, so we feed an
        # integer-compatible id from the stubbed run info.
        run_info.experiment_id = 4242
        yield run_info

    mlflow_repository_stub.start_run.side_effect = fake_start_run
    mlflow_repository_stub.log_params = MagicMock(return_value=None)
    mlflow_repository_stub._client.search_model_versions.return_value = [new_version]
    mlflow_repository_stub.promote_to_alias = MagicMock(return_value=None)


@pytest.mark.asyncio
@pytest.mark.integration
async def test_minimal_retrain_happy_path_writes_success_report(
    temporal_test_env: WorkflowEnvironment,
    temporal_worker_minimal_retrain: Worker,
    test_activities: Activities,
    postgres_engine,
    mlflow_repository_stub,
):
    """
    Scenario MR.1.1: Retrain succeeds. ``sientia_data.log_retrain`` must
    contain a success row with version/mlflow_run_id/mlflow_experiment_id
    populated and the registry must have been told to promote the new version
    to the configured alias.
    """
    client = temporal_test_env.client
    model_id = 711

    _seed_retrain_training_rows(postgres_engine, model_id)
    _configure_retrain_happy_path(mlflow_repository_stub)

    input_data = _retrain_input(model_id)

    with patch('laborious.activities.mlflow.mlflow.log_artifact') as log_artifact_mock:
        await start_and_await_workflow(
            client,
            MinimalRetrain.run,
            input_data,
            make_workflow_id('test-retrain-happy'),
        )

    assert log_artifact_mock.called, 'retrain_model should log the input CSV artifact'

    with postgres_engine.connect() as conn:
        rows = (
            conn.execute(
                text(
                    'SELECT * FROM sientia_data.log_retrain '
                    'WHERE model_id = :m'
                ),
                {'m': str(model_id)},
            )
            .mappings()
            .all()
        )

    assert len(rows) == 1
    for column in EXPECTED_RETRAIN_REPORT_COLUMNS:
        assert column in rows[0], f'Missing log_retrain column: {column}'

    row = rows[0]
    # ``model_id`` is now ``text``; compare against the stringified id.
    assert row['model_id'] == str(model_id)
    assert row['model_name'] == 'test_model'
    assert row['status'] == 'Model retrained successfully.'
    assert row['version'] == '7'
    assert row['mlflow_run_id'] == 'retrain-run-id'
    # ``mlflow_experiment_id`` is now ``int8``; assert the integer value
    # provided by the stubbed run info.
    assert row['mlflow_experiment_id'] == 4242
    assert row['timestamp'] is not None

    mlflow_repository_stub.promote_to_alias.assert_called_once()
    promote_kwargs = mlflow_repository_stub.promote_to_alias.call_args.kwargs
    assert promote_kwargs['model_name'] == 'test_model'
    assert promote_kwargs['version'] == '7'
    assert promote_kwargs['alias'] == 'production'


@pytest.mark.asyncio
@pytest.mark.integration
async def test_minimal_retrain_failure_writes_report_without_version_columns(
    temporal_test_env: WorkflowEnvironment,
    temporal_worker_minimal_retrain: Worker,
    test_activities: Activities,
    postgres_engine,
    mlflow_repository_stub,
):
    """
    Scenario MR.2.1: ``wrapper.retrain`` raises. The activity must catch the
    error, return ``success=False`` so ``update_production_model`` is skipped,
    and ``format_retrain_report`` must produce a row with the error message
    and NULL version columns.
    """
    client = temporal_test_env.client
    model_id = 721

    _seed_retrain_training_rows(postgres_engine, model_id)
    _configure_retrain_happy_path(mlflow_repository_stub)
    mlflow_repository_stub.get_cached_model.return_value.retrain.side_effect = RuntimeError(
        'training did not converge'
    )

    input_data = _retrain_input(model_id)

    with patch('laborious.activities.mlflow.mlflow.log_artifact'):
        await start_and_await_workflow(
            client,
            MinimalRetrain.run,
            input_data,
            make_workflow_id('test-retrain-failure'),
        )

    with postgres_engine.connect() as conn:
        rows = (
            conn.execute(
                text(
                    'SELECT * FROM sientia_data.log_retrain '
                    'WHERE model_id = :m'
                ),
                {'m': str(model_id)},
            )
            .mappings()
            .all()
        )

    assert len(rows) == 1
    row = rows[0]
    assert row['model_id'] == str(model_id)
    assert row['model_name'] == 'test_model'
    assert 'training did not converge' in row['status']
    assert row['version'] is None
    assert row['mlflow_run_id'] is None
    assert row['mlflow_experiment_id'] is None

    mlflow_repository_stub.promote_to_alias.assert_not_called()


@pytest.mark.asyncio
@pytest.mark.integration
async def test_minimal_retrain_missing_target_writes_failure_report(
    temporal_test_env: WorkflowEnvironment,
    temporal_worker_minimal_retrain: Worker,
    test_activities: Activities,
    postgres_engine,
    mlflow_repository_stub,
):
    """
    Scenario MR.2.2: ``model_config`` does not declare ``target``. The retrain
    activity must short-circuit before any MLflow call and the report row must
    carry the explicit guard message.
    """
    client = temporal_test_env.client
    model_id = 722

    _seed_retrain_training_rows(postgres_engine, model_id)
    _configure_retrain_happy_path(mlflow_repository_stub)

    input_data = _retrain_input(model_id, model_config={})

    with patch('laborious.activities.mlflow.mlflow.log_artifact'):
        await start_and_await_workflow(
            client,
            MinimalRetrain.run,
            input_data,
            make_workflow_id('test-retrain-missing-target'),
        )

    with postgres_engine.connect() as conn:
        row = (
            conn.execute(
                text(
                    'SELECT * FROM sientia_data.log_retrain '
                    'WHERE model_id = :m'
                ),
                {'m': str(model_id)},
            )
            .mappings()
            .first()
        )

    assert row is not None
    assert 'target' in row['status'].lower(), (
        f"expected target-missing message, got status={row['status']!r}"
    )
    assert row['version'] is None
    mlflow_repository_stub.get_cached_model.assert_not_called()


@pytest.mark.asyncio
@pytest.mark.integration
async def test_minimal_retrain_no_training_data_does_not_persist_report(
    temporal_test_env: WorkflowEnvironment,
    temporal_worker_minimal_retrain: Worker,
    test_activities: Activities,
    postgres_engine,
    mlflow_repository_stub,
):
    """
    Scenario MR.3.1: When the training query returns no rows the workflow must
    not persist any report row. The workflow currently raises plain
    ``ValueError`` which Temporal treats as a workflow-task failure (causing
    indefinite retries until the test environment times out), so the assertion
    here is constrained to the persistence side-effect. See ``e2e/CODE_ISSUES.md``
    issue MR-1 for the recommended ``ApplicationError`` fix.
    """
    client = temporal_test_env.client
    model_id = 731

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

    _configure_retrain_happy_path(mlflow_repository_stub)

    input_data = _retrain_input(model_id)

    with pytest.raises(Exception), patch('laborious.activities.mlflow.mlflow.log_artifact'):
        await start_and_await_workflow(
            client,
            MinimalRetrain.run,
            input_data,
            make_workflow_id('test-retrain-no-data'),
        )

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