"""
End-to-end test for retraining a model served through the explicit `joblib` flavor.

Retraining is not part of the `predictions_batch` workflow, so it gets its own module. The test
is a write→read round trip in one run: `retrain_model` loads the production artifact through the
joblib flavor, `log_model` writes the retrained model back as `model.pkl` + `code/`, and the
written artifact is then loaded again through the same flavor. Neither `joblib.dump` nor
`joblib.load` is mocked; only MLflow's own I/O is.
"""

import importlib
import shutil
import sys
from pathlib import Path

import pandas as pd
import pytest

from e2e.conftest import JOBLIB_FLAVOR_CODE_MODULE, JOBLIB_FLAVOR_MODEL_NAME
from laborious.activities.activities import Activities

RETRAIN_INDEX = [
    '2024-01-01 12:00:00+0000',
    '2024-01-01 12:01:00+0000',
    '2024-01-01 12:02:00+0000',
]

METADATA = {
    'model_id': 'test_model',
    'model_name': JOBLIB_FLAVOR_MODEL_NAME,
    'workflow_name': 'test_workflow',
    'schema_name': 'test_schedule',
}


def build_retrain_data() -> pd.DataFrame:
    """Training frame indexed by timestamp strings, carrying the target column."""
    return pd.DataFrame(
        {
            'sensor_1': [60.0, 61.0, 62.0],
            'sensor_2': [78.2, 78.4, 78.6],
            'target': [1.0, 2.0, 3.0],
        },
        index=RETRAIN_INDEX,
    )


class FakeArtifactStore:
    """Minimal stand-in for the run's artifact tree, so the written layout is inspectable."""

    def __init__(self, root: Path):
        self.root = root

    def log_artifact(self, local_path, artifact_path=None):
        if artifact_path is None:
            # Retraining also logs the CSVs at the run root; irrelevant to this test.
            return
        target = self.root / artifact_path
        target.mkdir(parents=True, exist_ok=True)
        shutil.copy(local_path, target / Path(local_path).name)

    def log_artifacts(self, local_dir, artifact_path=None):
        target = self.root / artifact_path if artifact_path else self.root
        shutil.copytree(local_dir, target, dirs_exist_ok=True)


@pytest.mark.asyncio
@pytest.mark.integration
async def test_retrain_joblib_model_writes_pickle_and_code_and_reloads(
    test_activities: Activities,
    joblib_flavor_artifacts,
    mlflow_client,
    patch_mlflow,
    tmp_path,
):
    """A joblib retrain writes `model.pkl` + `code/`, and the joblib flavor reads it back."""
    repository = test_activities.model_monitoring_repository

    store = FakeArtifactStore(tmp_path / 'run-retrained')
    patch_mlflow.log_artifact.side_effect = store.log_artifact
    patch_mlflow.log_artifacts.side_effect = store.log_artifacts
    patch_mlflow.get_run.return_value.data.params = {}
    patch_mlflow.start_run.return_value.__enter__.return_value.info.run_id = 'run-retrained'

    output = await repository.retrain_model(
        data=build_retrain_data(),
        model_name=JOBLIB_FLAVOR_MODEL_NAME,
        model_config={
            'target': 'target',
            'transform_flavor': 'joblib',
            'predict_flavor': 'joblib',
        },
        metadata=METADATA,
    )

    assert output['success'] is True, output.get('traceback')

    # Both models were loaded through the joblib flavor, not through a flavor loader.
    downloaded = [loaded.args[1] for loaded in mlflow_client.download_artifacts.call_args_list]
    assert sorted(downloaded) == ['data_model', 'prediction_model']
    assert patch_mlflow.sklearn.load_model.call_count == 0
    assert patch_mlflow.pyfunc.load_model.call_count == 0

    written = store.root / 'prediction_model'
    assert (written / 'model.pkl').is_file()
    assert (written / 'code' / f'{JOBLIB_FLAVOR_CODE_MODULE}.py').is_file()
    assert not (written / 'MLmodel').exists()
    assert (store.root / 'data_model' / 'model.pkl').is_file()

    # Drop the module the first load put in place: the reload must resolve it from the `code/`
    # the retrain wrote, not from what is already imported.
    for code_dir in (artifact['code_dir'] for artifact in joblib_flavor_artifacts.values()):
        if code_dir in sys.path:
            sys.path.remove(code_dir)
    sys.modules.pop(JOBLIB_FLAVOR_CODE_MODULE, None)
    importlib.invalidate_caches()

    mlflow_client.download_artifacts.side_effect = lambda _run_id, _artifact_name: str(written)

    reloaded = await repository.load_predict_model(JOBLIB_FLAVOR_MODEL_NAME, METADATA, 'joblib')

    features = pd.DataFrame({'feature_1': [0.234], 'feature_2': [0.783]})
    expected = joblib_flavor_artifacts['prediction_model']['model'].predict(features)

    assert reloaded.predict(features) == pytest.approx(expected)
    assert str(written / 'code') in sys.path