import json

import pytest

from mlflow.entities import (
    Dataset,
    DatasetInput,
    LifecycleStage,
    LoggedModelOutput,
    Metric,
    Run,
    RunData,
    RunInfo,
    RunInputs,
    RunOutputs,
    RunStatus,
)
from mlflow.exceptions import MlflowException

from tests.entities.test_run_data import _check as run_data_check
from tests.entities.test_run_info import _check as run_info_check
from tests.entities.test_run_inputs import _check as run_inputs_check


def _check_run(run, ri, rd_metrics, rd_params, rd_tags, datasets):
    run_info_check(
        run.info,
        ri.run_id,
        ri.experiment_id,
        ri.user_id,
        ri.status,
        ri.start_time,
        ri.end_time,
        ri.lifecycle_stage,
        ri.artifact_uri,
    )
    run_data_check(run.data, rd_metrics, rd_params, rd_tags)
    run_inputs_check(run.inputs, datasets)


def test_creation_and_hydration(run_data, run_info, run_inputs):
    run_data, metrics, params, tags = run_data
    (
        run_info,
        run_id,
        run_name,
        experiment_id,
        user_id,
        status,
        start_time,
        end_time,
        lifecycle_stage,
        artifact_uri,
    ) = run_info
    run_inputs, datasets = run_inputs
    run_outputs = RunOutputs(model_outputs=[LoggedModelOutput(model_id="model-id-1", step=3)])

    run1 = Run(run_info, run_data, run_inputs, run_outputs)

    _check_run(run1, run_info, metrics, params, tags, datasets)

    expected_info_dict = {
        "run_id": run_id,
        "run_name": run_name,
        "experiment_id": experiment_id,
        "user_id": user_id,
        "status": status,
        "start_time": start_time,
        "end_time": end_time,
        "lifecycle_stage": lifecycle_stage,
        "artifact_uri": artifact_uri,
    }
    assert run1.to_dictionary() == {
        "info": expected_info_dict,
        "data": {
            "metrics": {m.key: m.value for m in metrics},
            "params": {p.key: p.value for p in params},
            "tags": {t.key: t.value for t in tags},
        },
        "inputs": {
            "dataset_inputs": [
                {
                    "dataset": {
                        "digest": "digest1",
                        "name": "name1",
                        "profile": None,
                        "schema": None,
                        "source": "source",
                        "source_type": "my_source_type",
                    },
                    "tags": {"key": "value"},
                }
            ],
            "model_inputs": [],
        },
        "outputs": {
            "model_outputs": [{"model_id": "model-id-1", "step": 3}],
        },
    }
    # Run must be json serializable
    json.dumps(run1.to_dictionary())

    proto = run1.to_proto()
    run2 = Run.from_proto(proto)
    _check_run(run2, run_info, metrics, params, tags, datasets)
    assert run2.outputs.model_outputs == [LoggedModelOutput(model_id="model-id-1", step=3)]
    assert run2.outputs.to_dictionary() == {
        "model_outputs": [{"model_id": "model-id-1", "step": 3}],
    }

    run3 = Run(run_info, None, None)
    assert run3.to_dictionary() == {"info": expected_info_dict}

    run4 = Run(run_info, None)
    assert run4.to_dictionary() == {"info": expected_info_dict}


def test_string_repr():
    run_info = RunInfo(
        run_id="hi",
        run_name="name",
        experiment_id=0,
        user_id="user-id",
        status=RunStatus.FAILED,
        start_time=0,
        end_time=1,
        lifecycle_stage=LifecycleStage.ACTIVE,
    )
    metrics = [Metric(key=f"key-{i}", value=i, timestamp=0, step=i) for i in range(3)]
    run_data = RunData(metrics=metrics, params=[], tags=[])
    dataset_inputs = DatasetInput(
        dataset=Dataset(
            name="name1", digest="digest1", source_type="my_source_type", source="source"
        ),
        tags=[],
    )
    run_inputs = RunInputs(dataset_inputs=dataset_inputs)
    run1 = Run(run_info, run_data, run_inputs)
    expected = (
        "<Run: data=<RunData: metrics={'key-0': 0, 'key-1': 1, 'key-2': 2}, "
        "params={}, tags={}>, info=<RunInfo: artifact_uri=None, end_time=1, "
        "experiment_id=0, lifecycle_stage='active', run_id='hi', run_name='name', "
        "start_time=0, status=4, user_id='user-id'>, inputs=<RunInputs: "
        "dataset_inputs=<DatasetInput: dataset=<Dataset: digest='digest1', "
        "name='name1', profile=None, schema=None, source='source', "
        "source_type='my_source_type'>, tags=[]>, model_inputs=[]>, outputs=None>"
    )
    assert str(run1) == expected


def test_creating_run_with_absent_info_throws_exception(run_data, run_inputs):
    run_data = run_data[0]
    with pytest.raises(MlflowException, match="run_info cannot be None"):
        Run(None, run_data, run_inputs)