"""
End-to-end tests for the explicit `joblib` flavor on the MLflow model load paths.

`joblib.load` is NOT mocked here: the artifact fixture holds a real object serialized with
`joblib.dump(..., compress=3)` whose class lives only inside the artifact's `code/` directory.
A passing run therefore proves the flavor decompresses the pickle **and** resolves the module,
without any `mlflow.*.load_model` being called for that model.
"""

from unittest.mock import MagicMock

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

from e2e.conftest import JOBLIB_FLAVOR_MODEL_NAME
from e2e.helpers import (
    assert_prediction,
    assert_stop,
    insert_sample_data,
    make_workflow_id,
    start_and_await_workflow,
)
from laborious.activities.activities import Activities
from laborious.workflows.predictions_batch import PredictionsBatch

base_query = 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = {model_id}'

FLAVOR_LOADERS = ('sklearn', 'pyfunc')


def build_input_data(
    model_id: int,
    model_name: str,
    transform_flavor: str = 'sklearn',
    predict_flavor: str = 'sklearn',
) -> dict:
    """Workflow input for a plain transform + predict run of a given model."""
    return {
        'schedule_name': 'test-schedule',
        'model_name': model_name,
        'model_id': model_id,
        'query': base_query.format(model_id=model_id),
        'schema': 'predictions_schema',
        'table_name': 'predictions',
        'transform_table_name': 'transformed_data',
        'input_filters': {
            'SPECIFIC_VARIABLES_NULL_VALUES': {
                'POLICY': 'CONTINUE',
                'CONFIG': {'variables': ['sensor_1']},
            },
        },
        'mlflow_transform_filters': {
            'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}},
        },
        'mlflow_predict_filters': {
            'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}},
        },
        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
        'opc_output_config': {},
        'pi_web_api_output_config': {},
        'save_transform': True,
        'prediction_store_policy': 'lts:1',
        'model_config': {
            'retention_minutes': 0,
            'transform_flavor': transform_flavor,
            'predict_flavor': predict_flavor,
        },
        'datetime_columns': ['timestamp', 'created_at'],
    }


def expected_joblib_prediction(model) -> float:
    """Prediction the artifact's model produces for the transform output of this suite."""
    features = pd.DataFrame(
        {
            'feature_1': [0.234],
            'feature_2': [0.783],
            'timestamp': ['2024-01-01 12:00:00'],
        }
    )
    return float(model.predict(features)[0])


def assert_single_prediction_close(postgres_engine, model_id, expected, tolerance=1e-9):
    """Assert exactly one Good prediction row whose value matches `expected`."""
    with postgres_engine.connect() as conn:
        rows = conn.execute(
            text(
                'SELECT prediction, prediction_status '
                'FROM predictions_schema.predictions '
                f'WHERE model_id = {model_id}'
            )
        ).fetchall()

    assert len(rows) == 1, f'Expected one prediction record, got {len(rows)}'
    assert abs(float(rows[0][0]) - expected) < tolerance, (
        f'Expected prediction close to {expected}, got {rows[0][0]}'
    )
    assert rows[0][1] == 'Good', f"Expected prediction_status='Good', got {rows[0][1]}"


def logged_messages(mock_logger, level: str) -> str:
    """Concatenate every message logged at a given custom level."""
    calls = getattr(mock_logger, f'custom_{level}').call_args_list
    return ' '.join(str(logged.args[0]) for logged in calls)


def loader_uris(patch_mlflow) -> list[str]:
    """Every model URI handed to an MLflow flavor loader during the run."""
    return [
        str(loaded.args[0])
        for flavor in FLAVOR_LOADERS
        for loaded in getattr(patch_mlflow, flavor).load_model.call_args_list
    ]


@pytest.mark.asyncio
@pytest.mark.integration
async def test_prediction_workflow_with_joblib_flavor(
    temporal_test_env: WorkflowEnvironment,
    temporal_worker: Worker,
    test_activities: Activities,
    postgres_engine,
    joblib_flavor_artifacts,
    mlflow_client,
    patch_mlflow,
):
    """`predict_flavor: 'joblib'` deserializes the real artifact, and the prediction lands."""
    client = temporal_test_env.client
    model_id = 271
    insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
    input_data = build_input_data(model_id, JOBLIB_FLAVOR_MODEL_NAME, predict_flavor='joblib')

    await start_and_await_workflow(
        client, PredictionsBatch.run, input_data, make_workflow_id('test-joblib-flavor')
    )

    assert_single_prediction_close(
        postgres_engine,
        model_id,
        expected_joblib_prediction(joblib_flavor_artifacts['prediction_model']['model']),
    )

    assert mlflow_client.download_artifacts.call_count == 1
    assert mlflow_client.download_artifacts.call_args.args[1] == 'prediction_model'

    # The joblib flavor never builds a `models:/` URI, so no flavor loader saw this model.
    assert all(JOBLIB_FLAVOR_MODEL_NAME not in uri for uri in loader_uris(patch_mlflow))


@pytest.mark.asyncio
@pytest.mark.integration
async def test_prediction_workflow_with_joblib_transform_flavor(
    temporal_test_env: WorkflowEnvironment,
    temporal_worker: Worker,
    test_activities: Activities,
    postgres_engine,
    joblib_flavor_artifacts,
    mlflow_client,
    patch_mlflow,
    mock_mlflow_models,
):
    """`transform_flavor: 'joblib'` downloads and deserializes `data_model` on its own."""
    client = temporal_test_env.client
    model_id = 274
    insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
    input_data = build_input_data(model_id, JOBLIB_FLAVOR_MODEL_NAME, transform_flavor='joblib')

    await start_and_await_workflow(
        client, PredictionsBatch.run, input_data, make_workflow_id('test-joblib-transform')
    )

    assert_prediction(postgres_engine, model_id, prediction=0.5)
    mock_mlflow_models['predict_model'].predict.assert_called()

    assert mlflow_client.download_artifacts.call_count == 1
    assert mlflow_client.download_artifacts.call_args.args[1] == 'data_model'

    # The joblib flavor resolves no artifact URI, so `get_run` is never reached and no flavor
    # loader was given a `data_model` URI.
    patch_mlflow.get_run.assert_not_called()
    assert all(not uri.endswith('/data_model') for uri in loader_uris(patch_mlflow))


@pytest.mark.asyncio
@pytest.mark.integration
async def test_prediction_workflow_with_sklearn_flavor_is_unchanged(
    temporal_test_env: WorkflowEnvironment,
    temporal_worker: Worker,
    test_activities: Activities,
    postgres_engine,
    mlflow_client,
    mock_mlflow_models,
):
    """A model on an MLflow flavor keeps the current result and downloads nothing."""
    client = temporal_test_env.client
    model_id = 272
    insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
    input_data = build_input_data(model_id, 'test_model')

    await start_and_await_workflow(
        client, PredictionsBatch.run, input_data, make_workflow_id('test-joblib-sklearn-flavor')
    )

    assert_prediction(postgres_engine, model_id, prediction=0.5)
    mock_mlflow_models['predict_model'].predict.assert_called()
    mlflow_client.download_artifacts.assert_not_called()


@pytest.mark.asyncio
@pytest.mark.integration
async def test_prediction_workflow_downloads_only_the_artifact_code_directory(
    temporal_test_env: WorkflowEnvironment,
    temporal_worker: Worker,
    test_activities: Activities,
    postgres_engine,
    joblib_flavor_artifacts,
    mlflow_client,
    mock_mlflow_models,
):
    """A flavor model whose artifact carries `code` downloads that directory — and only it.

    The flavor loaders read the model URI themselves, so nothing else is fetched through the
    client: the `code` directory is the whole of what this workflow downloads.
    """
    mlflow_client.list_artifacts.side_effect = lambda _run_id, artifact_path=None: (
        [MagicMock(path=f'{artifact_path}/code')] if artifact_path else []
    )

    client = temporal_test_env.client
    model_id = 275
    insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
    input_data = build_input_data(model_id, 'test_model')

    await start_and_await_workflow(
        client, PredictionsBatch.run, input_data, make_workflow_id('test-joblib-code-only')
    )

    assert_prediction(postgres_engine, model_id, prediction=0.5)

    downloaded = [loaded.args[1] for loaded in mlflow_client.download_artifacts.call_args_list]
    assert sorted(downloaded) == ['data_model/code', 'prediction_model/code']


@pytest.mark.asyncio
@pytest.mark.integration
async def test_prediction_workflow_joblib_flavor_failure_reports_single_cause(
    temporal_test_env: WorkflowEnvironment,
    temporal_worker: Worker,
    test_activities: Activities,
    postgres_engine,
    joblib_flavor_artifact_without_pickle,
    mlflow_client,
    mock_logger,
):
    """The joblib flavor failing stops the workflow with one cause, not a chain of two."""
    mlflow_client.download_artifacts.side_effect = lambda _run_id, _artifact_name: (
        joblib_flavor_artifact_without_pickle
    )

    client = temporal_test_env.client
    model_id = 273
    insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
    input_data = build_input_data(model_id, JOBLIB_FLAVOR_MODEL_NAME, predict_flavor='joblib')

    await start_and_await_workflow(
        client, PredictionsBatch.run, input_data, make_workflow_id('test-joblib-flavor-failed')
    )

    assert_stop(postgres_engine, model_id)

    logged = logged_messages(mock_logger, 'debug')
    assert 'No .pkl file found in prediction_model artifact' in logged
    assert 'Original error' not in logged