import mlflow
from mlflow.tracking import MlflowClient

def register_and_promote_model(
    model_uri: str,
    model_name: str,
    val_f1: float,
    robustness_drop: float = 0.0,
    description: str = "",
    tags: dict = None
) -> str:
    """
    Registers model in the MLflow Model Registry and promotes alias to 'champion'
    if it surpasses previous performance or if no champion exists.
    """
    client = MlflowClient()

    # Ensure registered model exists
    try:
        client.create_registered_model(model_name)
    except Exception:
        pass

    # Create model version
    mv = client.create_model_version(
        name=model_name,
        source=model_uri,
        run_id=model_uri.split("/")[1] if "runs:/" in model_uri else "",
        description=description,
        tags=tags or {}
    )

    version = mv.version

    # Check existing champion
    try:
        champ_version = client.get_model_version_by_alias(model_name, "champion")
        prev_f1 = float(champ_version.tags.get("val_f1", 0.0))
        if val_f1 > prev_f1 and robustness_drop < 0.25:
            client.set_registered_model_alias(model_name, "champion", version)
            client.set_registered_model_alias(model_name, "challenger", champ_version.version)
            action = "PROMOTED_CHAMPION"
        else:
            client.set_registered_model_alias(model_name, "challenger", version)
            action = "SET_CHALLENGER"
    except Exception:
        # No champion yet, promote current
        client.set_registered_model_alias(model_name, "champion", version)
        action = "FIRST_CHAMPION"

    # Save validation metrics as tags
    client.set_model_version_tag(model_name, version, "val_f1", str(val_f1))
    client.set_model_version_tag(model_name, version, "robustness_drop", str(robustness_drop))

    return action

def get_champion_model_uri(model_name: str) -> str:
    """Returns the MLflow models URI for the current champion alias."""
    return f"models:/{model_name}@champion"