"""
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
import pandas as pd
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
from sientia_model.wrappers.sientia_model import SientiaModel

# 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',
]

# Matches ``retrain_model`` return ``message`` when ``success`` is True (also written to ``log_retrain.status``).
RETRAIN_ACTIVITY_SUCCESS_MESSAGE = 'Model retrained successfully.'


class _FakeSientiaModelForMinimalRetrain(SientiaModel):
    """
    Fake SientiaModel that uses the real SientiaModel lifecycle to surface
    index-alignment issues during ``retrain()``.

    It intentionally performs strict alignment inside ``_retrain_model``:
    ``y.loc[x.index]``.
    """

    def __init__(self, *, target: str = 'sensor_1'):
        super().__init__(
            model_type='FakeMinimalRetrain',
            model_version='0.0.0',
            model=object(),
            transformer=object(),
        )
        self.target = target
        self.model_is_fitted = True
        self.force_retrain_error = False

    def store_model(  # type: ignore[override]
        self,
        name: str,
        signature=None,
        pip_requirements=None,
        code_path=None,
    ) -> None:
        # No-op: E2E tests validate workflow persistence, not real MLflow artifacts.
        return None

    def _predict(self, data: pd.DataFrame):
        pred = pd.DataFrame({'prediction': [0.5] * len(data)}, index=data.index)
        return pred, {}

    def _transform(self, data: pd.DataFrame):
        out = data.drop(columns=[self.target], errors='ignore').copy()
        out.index = data.index
        return out, {}

    def _train_transformer(self, train_data: pd.DataFrame, val_data: pd.DataFrame) -> None:
        return None

    def _train_model(
        self,
        x: pd.DataFrame,
        y: pd.DataFrame,
        x_val: pd.DataFrame | None = None,
        y_val: pd.DataFrame | None = None,
    ) -> None:
        return None

    def _retrain_transformer(self, data: pd.DataFrame) -> None:
        return None

    def _retrain_model(self, x: pd.DataFrame, y: pd.DataFrame | None) -> None:
        if self.force_retrain_error:
            raise RuntimeError('training did not converge')
        if y is None:
            return
        # Strict alignment on purpose to reproduce the production failure mode.
        _ = y.loc[x.index]


@pytest.fixture
def mlflow_repository_stub():
    """
    Override the shared E2E fixture: return a real fake ``SientiaModel`` wrapper
    instead of a MagicMock wrapper.
    """
    repo = MagicMock()
    repo._client = MagicMock()

    wrapper = _FakeSientiaModelForMinimalRetrain(target='sensor_1')
    repo.get_cached_model = MagicMock(return_value=wrapper)
    return repo


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``).
    - ``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

    @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'),
        )

    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]
    for column in EXPECTED_RETRAIN_REPORT_COLUMNS:
        assert column in row, f'Missing log_retrain column: {column}'

    assert row['status'] == RETRAIN_ACTIVITY_SUCCESS_MESSAGE, (
        "Expected retrain_model to return success (experiment_response['success'] is True). "
        'Persisted log_retrain.status is the activity message; when success is False the run '
        'never reaches mlflow.log_artifact — diagnose the retrain failure from status below, '
        'not from a skipped artifact upload. '
        f"Got status={row['status']!r}, version={row.get('version')!r}, "
        f"mlflow_run_id={row.get('mlflow_run_id')!r}."
    )

    assert log_artifact_mock.called, (
        'After a successful retrain, retrain_model must call mlflow.log_artifact for the '
        'input CSV inside start_run.'
    )

    # ``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['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.force_retrain_error = True

    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