from __future__ import annotations

import os
import subprocess
import sys
import textwrap
from pathlib import Path

from mlflow.pytest.session import TAG_TEST_NAME
from mlflow.tracking import MlflowClient

_GENERATED_TEST = """
import mlflow
from mlflow.genai.scorers import scorer


@scorer
def always_pass(*, outputs):
    return True


@mlflow.test
def test_marked():
    result = mlflow.genai.evaluate(
        data=[{"inputs": {"text": "hi"}, "outputs": "hi"}],
        scorers=[always_pass],
    )
    assert result.passed, result.reason
"""


def test_evaluate_traces_tagged_with_test_identity(tmp_path: Path):
    test_file = tmp_path / "test_generated.py"
    test_file.write_text(textwrap.dedent(_GENERATED_TEST))

    tracking_uri = f"sqlite:///{tmp_path / 'mlflow.db'}"
    result = subprocess.run(
        [
            sys.executable,
            "-m",
            "pytest",
            test_file.name,
            "-p",
            "no:cacheprovider",
            # The plugin is opt-in (no pytest11 entry point).
            "-p",
            "mlflow.pytest.plugin",
            "-q",
        ],
        cwd=tmp_path,
        capture_output=True,
        text=True,
        env={**os.environ, "MLFLOW_TRACKING_URI": tracking_uri},
    )
    assert result.returncode == 0, result.stdout + result.stderr

    client = MlflowClient(tracking_uri=tracking_uri)
    experiment_ids = [e.experiment_id for e in client.search_experiments()]
    traces = client.search_traces(experiment_ids=experiment_ids) if experiment_ids else []
    assert traces
    assert all(t.info.tags.get(TAG_TEST_NAME, "").endswith("::test_marked") for t in traces)