from azure.identity import ClientSecretCredential
from azure.keyvault.secrets import SecretClient
from azure.ai.ml import MLClient
from azure.ai.ml.entities import Model

# Static config
key_vault_url = "https://donkv.vault.azure.net/"

# TEMPORARY bootstrap credential - this comes from the Azure ML job identity
bootstrap_cred = ClientSecretCredential(
    tenant_id="6ad27142-5e7d-4841-9e72-c3aaca00f9e6",
    client_id="a3b5f924-df36-4138-bd7e-f9d85544af10",
    client_secret="-FF8Q~XqBZGpNN_0Y_MPHbWVHTM5rqGzJvZXTbai"
)

# Connect to Key Vault and retrieve secrets (optionally re-use bootstrap credentials)
secret_client = SecretClient(vault_url=key_vault_url, credential=bootstrap_cred)

client_id = secret_client.get_secret("AZURE-CLIENT-ID").value
tenant_id = secret_client.get_secret("AZURE-TENANT-ID").value
client_secret = secret_client.get_secret("AZURE-CLIENT-SECRET").value
subscription_id = secret_client.get_secret("AZURE-SUBSCRIPTION-ID").value
resource_group = secret_client.get_secret("AZURE-RESOURCE-GROUP").value
workspace_name = secret_client.get_secret("AZURE-WORKSPACE-NAME").value

# Use retrieved secrets to authenticate with Azure ML
ml_cred = ClientSecretCredential(tenant_id=tenant_id, client_id=client_id, client_secret=client_secret)
ml_client = MLClient(ml_cred, subscription_id, resource_group, workspace_name)

# Register the model
model = Model(
    path="outputs/model.joblib",
    name="iris-logreg",
    description="Logistic regression model trained on iris dataset",
    type="custom_model"
)
registered_model = ml_client.models.create_or_update(model)
print(f"✅ Model registered: {registered_model.name}, version: {registered_model.version}")