"""Verify that every YOLO smoke task was persisted completely in MLflow."""

from __future__ import annotations

import argparse
import json

import mlflow
from mlflow.entities import Run
from mlflow.tracking import MlflowClient


TASKS = ("detect", "segment", "classify", "pose", "obb")
REQUIRED_ARTIFACTS = {"weights/best.pt", "weights/last.pt", "results.csv"}


def artifact_paths(client: MlflowClient, run_id: str, path: str = "") -> set[str]:
    result: set[str] = set()
    for artifact in client.list_artifacts(run_id, path):
        if artifact.is_dir:
            result.update(artifact_paths(client, run_id, artifact.path))
        else:
            result.add(artifact.path)
    return result


def latest_task_run(runs: list[Run], task: str) -> Run:
    expected_name = f"{task}-smoke"
    for run in runs:
        if run.data.tags.get("mlflow.runName") == expected_name:
            return run
    raise AssertionError(f"MLflow run not found: {expected_name}")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--tracking-uri",
        default="sqlite:///runs/yolo26_mlflow_smoke/mlflow.db",
    )
    parser.add_argument("--experiment", default="yolo26-mlflow-smoke")
    args = parser.parse_args()

    mlflow.set_tracking_uri(args.tracking_uri)
    client = MlflowClient()
    experiment = client.get_experiment_by_name(args.experiment)
    if experiment is None:
        raise AssertionError(f"MLflow experiment not found: {args.experiment}")

    runs = client.search_runs(
        [experiment.experiment_id],
        order_by=["start_time DESC"],
    )
    summary: dict[str, object] = {
        "tracking_uri": args.tracking_uri,
        "experiment_id": experiment.experiment_id,
        "artifact_location": experiment.artifact_location,
        "tasks": {},
    }
    task_summary: dict[str, object] = summary["tasks"]  # type: ignore[assignment]

    for task in TASKS:
        run = latest_task_run(runs, task)
        artifacts = artifact_paths(client, run.info.run_id)
        missing = REQUIRED_ARTIFACTS - artifacts
        assert run.info.status == "FINISHED", (task, run.info.status)
        assert run.data.params, f"No parameters logged for {task}"
        assert run.data.metrics, f"No metrics logged for {task}"
        assert not missing, f"Missing artifacts for {task}: {sorted(missing)}"
        task_summary[task] = {
            "run_id": run.info.run_id,
            "status": run.info.status,
            "parameters": len(run.data.params),
            "metrics": len(run.data.metrics),
            "artifact_uri": run.info.artifact_uri,
            "required_artifacts": sorted(REQUIRED_ARTIFACTS),
        }

    print(json.dumps(summary, indent=2, ensure_ascii=False))


if __name__ == "__main__":
    main()