"""An import that works, twice over: once from the producer's own bytes and once from the factory.

Two origins, because they prove different things. The **golden** bundle is bytes the producer's code
emitted — it decrypts at the producer's real Argon2id cost (256 MiB, 3 passes) and carries the tree
the producer really writes, so it is the only case that would notice the producer's layout moving.
The **factory** bundle is synthesised here, cheaply, with names unique to the test; it is what the
rest of the suite runs on, so a happy path over it is what makes the negative cases' "and nothing
was provisioned" meaningful.

What "it worked" means, in full: the experiment exists, a run exists inside it carrying the derived
run name, the artifact tree is under the run with `prediction_model` and `data_model` at its root,
every parameter the bundle declared is logged plus the origin's identity under `source.*`, a
registry version exists, the MongoDB document exists, and the status row ends `IMPORT_SUCCESS`.

Two things a successful import must **not** have done are asserted here too, because they are the
difference between "imported" and "live": no version is in `Production`, and the document is
`active: false`. Promotion and activation are manual operator steps (QTZPOC-21).
"""

from __future__ import annotations

import json
from typing import Any

import mlflow
import pytest

from e2e.import_harness import (
    IMPORT_BUCKET,
    IMPORT_PREFIX,
    SEEDED_MODEL_IDS,
)

pytestmark = [pytest.mark.import_e2e, pytest.mark.asyncio]

GOLDEN_MODEL_NAME = 'fixture_model'
GOLDEN_EXPERIMENT_NAME = 'fixture_experiment'
GOLDEN_MODEL_VERSION = '3'
GOLDEN_TARGET = 'SE'

REQUIRED_ARTIFACT_ROOTS = {'prediction_model', 'data_model'}


def _mlflow_client(server: dict[str, Any]) -> Any:
    return mlflow.tracking.MlflowClient(tracking_uri=server['tracking_uri'])


def _as_dict(value: Any) -> dict[str, Any]:
    """`request_data` / `orchestrator_response_data` come back as a dict or as JSON text."""
    return json.loads(value) if isinstance(value, str) else (value or {})


async def _run_import(
    *,
    payload: bytes,
    import_upload_bundle: Any,
    import_insert_status_row: Any,
    import_build_input: Any,
    import_execute: Any,
    password: str,
    import_envelope: Any,
) -> tuple[Any, int, str]:
    """Upload, insert the row the BFF would have inserted, run the workflow. Return the outcome."""
    object_key, digest = import_upload_bundle(payload)
    import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
    handle = await import_execute(
        import_build_input(
            import_run_id=import_run_id,
            object_key=object_key,
            expected_digest=digest,
            password_envelope=import_envelope(password),
        )
    )
    return handle, import_run_id, object_key


def _assert_provisioned(
    *,
    client: Any,
    experiment_name: str,
    model_name: str,
    run_name: str,
    expected_parameters: dict[str, str],
    expected_source_fields: dict[str, str],
) -> tuple[str, Any]:
    """The whole positive claim about MLflow, in one place. Returns `(run_id, version)`."""
    experiment = client.get_experiment_by_name(experiment_name)
    assert experiment is not None, f'no experiment named {experiment_name!r}'

    runs = [
        run
        for run in client.search_runs([experiment.experiment_id])
        if run.info.run_name == run_name
    ]
    assert len(runs) == 1, f'expected exactly one run named {run_name!r}, found {len(runs)}'
    run = runs[0]
    assert run.info.status == 'FINISHED'

    artifact_roots = {artifact.path for artifact in client.list_artifacts(run.info.run_id)}
    assert REQUIRED_ARTIFACT_ROOTS <= artifact_roots, (
        f'the run holds {sorted(artifact_roots)}, missing {sorted(REQUIRED_ARTIFACT_ROOTS - artifact_roots)}'
    )
    # `prediction_model` is a directory the legacy loader walks, not an opaque blob.
    prediction_files = {
        artifact.path.rsplit('/', 1)[-1]
        for artifact in client.list_artifacts(run.info.run_id, 'prediction_model')
    }
    assert {'MLmodel', 'model.pkl'} <= prediction_files, prediction_files

    for name, value in expected_parameters.items():
        assert run.data.params.get(name) == value, (
            f'parameter {name!r} is {run.data.params.get(name)!r}, expected {value!r}'
        )
    for name, value in expected_source_fields.items():
        assert run.data.params.get(f'source.{name}') == value, (
            f'source.{name} is {run.data.params.get(f"source.{name}")!r}, expected {value!r}'
        )

    versions = list(client.search_model_versions(f"name='{model_name}'"))
    assert len(versions) == 1, f'expected one registered version, found {len(versions)}'
    version = versions[0]
    assert version.run_id == run.info.run_id
    return run.info.run_id, version


class TestGoldenBundle:
    """The producer's own bytes, at the producer's own KDF cost."""

    async def test_the_producer_artefact_imports_end_to_end(
        self,
        import_status_table: Any,
        import_models_collection: Any,
        import_golden_bundle: dict[str, Any],
        import_mlflow_server: dict[str, Any],
        import_upload_bundle: Any,
        import_insert_status_row: Any,
        import_build_input: Any,
        import_execute: Any,
        import_envelope: Any,
        import_read_status_row: Any,
    ) -> None:
        handle, import_run_id, object_key = await _run_import(
            payload=import_golden_bundle['bytes'],
            import_upload_bundle=import_upload_bundle,
            import_insert_status_row=import_insert_status_row,
            import_build_input=import_build_input,
            import_execute=import_execute,
            password=import_golden_bundle['password'],
            import_envelope=import_envelope,
        )
        result = await handle.result()

        assert result['model_name'] == GOLDEN_MODEL_NAME
        assert result['import_run_id'] == import_run_id
        assert result['version'] == '1', 'the destination numbering starts at 1, not at the bundle'

        run_id, version = _assert_provisioned(
            client=_mlflow_client(import_mlflow_server),
            experiment_name=GOLDEN_EXPERIMENT_NAME,
            model_name=GOLDEN_MODEL_NAME,
            run_name=f'import-{GOLDEN_MODEL_NAME}-v{GOLDEN_MODEL_VERSION}',
            expected_parameters={'target_variable': GOLDEN_TARGET, 'nan_treatment': 'drop'},
            expected_source_fields={
                'model_name': GOLDEN_MODEL_NAME,
                'experiment_name': GOLDEN_EXPERIMENT_NAME,
                'model_version': GOLDEN_MODEL_VERSION,
                'model_project': 'fixture_project',
            },
        )
        assert result['run_id'] == run_id

        document = import_models_collection.find_one({'name': GOLDEN_MODEL_NAME}, {'_id': False})
        assert document is not None, 'no model document was written'
        assert document['id'] == str(max(int(i) for i in SEEDED_MODEL_IDS) + 1)
        assert document['model_config']['target'] == GOLDEN_TARGET

        row = import_read_status_row(import_run_id)
        assert row is not None
        assert row['status'] == 'IMPORT_SUCCESS'
        assert row['error_message'] is None
        assert row['experiment_name'] == GOLDEN_EXPERIMENT_NAME
        assert row['run_name'] == f'import-{GOLDEN_MODEL_NAME}-v{GOLDEN_MODEL_VERSION}'
        assert row['file_name'] == object_key
        assert row['bucket_name'] == IMPORT_BUCKET
        assert object_key.startswith(IMPORT_PREFIX)

        # Imported, not live.
        assert version.current_stage in (None, 'None'), version.current_stage
        assert document['active'] is False


class TestFactoryBundle:
    """A synthesised tree with names unique to this test — the shape the rest of the suite runs on."""

    async def test_a_factory_bundle_imports_end_to_end(
        self,
        import_status_table: Any,
        import_models_collection: Any,
        import_names: dict[str, str],
        import_bundle_factory: Any,
        import_mlflow_server: dict[str, Any],
        import_upload_bundle: Any,
        import_insert_status_row: Any,
        import_build_input: Any,
        import_execute: Any,
        import_envelope: Any,
        import_read_status_row: Any,
    ) -> None:
        payload = import_bundle_factory(parameters={'target_variable': 'SE'})
        handle, import_run_id, _ = await _run_import(
            payload=payload,
            import_upload_bundle=import_upload_bundle,
            import_insert_status_row=import_insert_status_row,
            import_build_input=import_build_input,
            import_execute=import_execute,
            password='Factory1Model!Test',
            import_envelope=import_envelope,
        )
        result = await handle.result()

        assert result['model_name'] == import_names['model_name']
        run_id, version = _assert_provisioned(
            client=_mlflow_client(import_mlflow_server),
            experiment_name=import_names['experiment_name'],
            model_name=import_names['model_name'],
            # DEFAULT_METADATA declares model_version '2'.
            run_name=f'import-{import_names["model_name"]}-v2',
            expected_parameters={'target_variable': 'SE', 'retrain': 'False'},
            expected_source_fields={
                'model_name': import_names['model_name'],
                'experiment_name': import_names['experiment_name'],
                'model_version': '2',
            },
        )
        assert result['run_id'] == run_id
        assert result['version'] == '1'

        document = import_models_collection.find_one(
            {'name': import_names['model_name']}, {'_id': False}
        )
        assert document is not None
        assert document['model_config'] == {
            'transform_flavor': 'sklearn',
            'predict_flavor': 'joblib',
            'retention_minutes': 60,
            'target': 'SE',
        }

        row = import_read_status_row(import_run_id)
        assert row is not None
        assert row['status'] == 'IMPORT_SUCCESS'
        detail = _as_dict(row['orchestrator_response_data'])
        assert detail['import']['workflow_id'] == handle.id

        assert version.current_stage in (None, 'None')
        assert document['active'] is False


class TestLegacyLoaderResolution:
    """The registry source an import writes has to be one the *existing* loader can parse.

    This is the assertion that could only ever be made against a real registry: the importer writes
    the run's resolved artifact URI with `prediction_model` appended, because `get_model_run_id`
    extracts the run id as `source.split('/')[2]` — an index that is the run id only in the resolved
    form. A literal `runs:/<run_id>/prediction_model` would put `prediction_model` at index 2, and a
    `models:/m-*` source would have none of the parts at all.

    It is called with the stage the version is actually in rather than `Production`, because an
    import deliberately promotes nothing — passing `Production` here would test the promotion step
    this feature does not have.
    """

    async def test_the_imported_source_resolves_to_the_import_run(
        self,
        import_status_table: Any,
        import_models_collection: Any,
        import_names: dict[str, str],
        import_bundle_factory: Any,
        import_activities: Any,
        import_mlflow_server: dict[str, Any],
        import_upload_bundle: Any,
        import_insert_status_row: Any,
        import_build_input: Any,
        import_execute: Any,
        import_envelope: Any,
    ) -> None:
        payload = import_bundle_factory()
        handle, _, _ = await _run_import(
            payload=payload,
            import_upload_bundle=import_upload_bundle,
            import_insert_status_row=import_insert_status_row,
            import_build_input=import_build_input,
            import_execute=import_execute,
            password='Factory1Model!Test',
            import_envelope=import_envelope,
        )
        result = await handle.result()
        model_name = import_names['model_name']

        client = _mlflow_client(import_mlflow_server)
        version = next(iter(client.search_model_versions(f"name='{model_name}'")))
        assert version.source.endswith('/prediction_model'), version.source

        # The production method, on the production repository object, against the real registry.
        resolved = import_activities.import_mlflow_repository.get_model_run_id(
            model_name, stage=version.current_stage
        )
        assert resolved == result['run_id'], (
            f'get_model_run_id resolved {resolved!r} from source {version.source!r}, '
            f'but the import created run {result["run_id"]!r}'
        )

        artifact_roots = {artifact.path for artifact in client.list_artifacts(resolved)}
        assert REQUIRED_ARTIFACT_ROOTS <= artifact_roots, sorted(artifact_roots)