"""Example 3: Model Registry Management.

Demonstrates how to use ModelRegistry to manage model versions,
promote models with aliases, and load models for inference.

Prerequisites:
    - A running MLflow tracking server
    - At least one model uploaded (see 02_upload_sklearn.py)

Run:
    python examples/03_registry_management.py
"""

from dgtw_mlflow_utils import ALIAS_CHAMPION, ModelRegistry


def main():
    """Main function to execute the example."""
    # ------------------------------------------------------------------
    # 1) Connect to the registry
    # ------------------------------------------------------------------
    registry = ModelRegistry(tracking_uri="http://127.0.0.1:5001")

    MODEL_NAME = "TAT_RandomForest_Baseline"  # from example 02

    # ------------------------------------------------------------------
    # 2) List all versions
    # ------------------------------------------------------------------
    print("=" * 60)
    print("  Listing Model Versions")
    print("=" * 60)

    versions = registry.list_versions(MODEL_NAME)

    if not versions:
        print(f"  No versions found for '{MODEL_NAME}'.")
        print("  Run 02_upload_sklearn.py first to upload a model.")
        return

    for v in versions:
        print(f"  Version {v.version}  |  Run ID: {v.run_id}")

    # ------------------------------------------------------------------
    # 3) Print a formatted summary
    # ------------------------------------------------------------------
    print()
    registry.print_summary(MODEL_NAME)

    # ------------------------------------------------------------------
    # 4) Promote latest version to "champion"
    # ------------------------------------------------------------------
    print("Promoting latest version to 'champion'...")
    promoted = registry.promote(MODEL_NAME, alias=ALIAS_CHAMPION)
    print(f"  ✓ Version {promoted.version} is now the champion\n")

    # ------------------------------------------------------------------
    # 5) Load the champion model for inference
    # ------------------------------------------------------------------
    print("Loading champion model...")
    model = registry.load_champion(MODEL_NAME)
    print(f"  ✓ Model loaded: {type(model)}\n")

    # Run a quick prediction (Iris features)
    import pandas as pd

    sample = pd.DataFrame(
        [[5.1, 3.5, 1.4, 0.2]],
        columns=["sepal length (cm)", "sepal width (cm)", "petal length (cm)", "petal width (cm)"],
    )
    prediction = model.predict(sample)
    print(f"  Sample prediction: {prediction}")

    # ------------------------------------------------------------------
    # 6) Tag a version
    # ------------------------------------------------------------------
    latest = registry.get_latest_version(MODEL_NAME)
    if latest:
        registry.set_tag(MODEL_NAME, latest.version, "approved_by", "ai-team")
        print(f"\n  ✓ Tagged version {latest.version} with approved_by=ai-team")

    # ------------------------------------------------------------------
    # 7) Cleanup old versions (dry run)
    # ------------------------------------------------------------------
    print("\nDry-run cleanup (keep 3 newest):")
    would_delete = registry.cleanup_old_versions(MODEL_NAME, keep=3, dry_run=True)
    if would_delete:
        print(f"  Would delete versions: {would_delete}")
    else:
        print("  Nothing to clean up.")

    print()


if __name__ == "__main__":
    main()
