"""Record a small calculation in an AI Experiments instance."""

import argparse
import json
import os

import mlflow

from get_token import get_token


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--numbers", nargs="+", type=int, default=[10, 20, 30])
    parser.add_argument("--experiment", default="first-experiment")
    parser.add_argument("--run-name", default="first-total")
    args = parser.parse_args()

    tracking_uri = os.environ.get("MLFLOW_TRACKING_URI", "").strip()
    if not tracking_uri:
        raise SystemExit("Set MLFLOW_TRACKING_URI to your instance URL first.")

    mlflow.set_tracking_uri(tracking_uri)
    os.environ["MLFLOW_TRACKING_TOKEN"] = get_token()
    mlflow.set_experiment(args.experiment)

    total = sum(args.numbers)
    with mlflow.start_run(run_name=args.run_name) as run:
        mlflow.log_param("numbers", json.dumps(args.numbers))
        mlflow.log_metric("total", total)
        mlflow.log_text(
            json.dumps({"numbers": args.numbers, "total": total}, indent=2),
            "result.json",
        )

    print(f"Total: {total}")
    print(f"Run ID: {run.info.run_id}")


if __name__ == "__main__":
    main()