"""
deploy_endpoint.py — Deploy a Model as a Managed Online Endpoint
================================================================
Creates a managed online endpoint in Azure AI Foundry with:
  - Endpoint creation
  - Model registration
  - Deployment with scaling configuration
  - Traffic routing
  - Health check

Prerequisites:
    pip install azure-ai-ml azure-identity python-dotenv rich

Usage:
    python deploy_endpoint.py
"""

import os
import sys
from dotenv import load_dotenv
from azure.identity import DefaultAzureCredential
from azure.ai.ml import MLClient
from azure.ai.ml.entities import (
    ManagedOnlineEndpoint,
    ManagedOnlineDeployment,
    Model,
    Environment,
    CodeConfiguration,
    OnlineRequestSettings,
    ProbeSettings,
)
from rich.console import Console

load_dotenv()
console = Console()


def get_env(key: str) -> str:
    value = os.getenv(key)
    if not value:
        console.print(f"[red]ERROR:[/red] '{key}' not set. See .env.example.")
        sys.exit(1)
    return value


def main() -> None:
    """Deploy a model as a managed online endpoint with scaling configuration."""

    # --- Authenticate and create ML client -----------------------------------
    credential = DefaultAzureCredential()
    ml_client = MLClient(
        credential=credential,
        subscription_id=get_env("AZURE_SUBSCRIPTION_ID"),
        resource_group_name=get_env("AZURE_RESOURCE_GROUP"),
        workspace_name=get_env("AZURE_AI_FOUNDRY_PROJECT_NAME"),
    )

    console.print("[bold]Managed Endpoint Deployment[/bold]", style="magenta")
    console.print("=" * 60)

    # --- Step 1: Create the endpoint -----------------------------------------
    endpoint_name = "ai-foundry-demo-endpoint"
    console.print(f"\n[bold]Step 1:[/bold] Creating endpoint '{endpoint_name}'...", style="cyan")

    endpoint = ManagedOnlineEndpoint(
        name=endpoint_name,
        description="Demo endpoint for Azure AI Foundry course",
        auth_mode="key",
        tags={"environment": "lab", "course": "ai-foundry"},
    )

    try:
        ml_client.online_endpoints.begin_create_or_update(endpoint).result()
        console.print(f"   ✅ Endpoint created: {endpoint_name}")
    except Exception as e:
        console.print(f"   ⚠️  Endpoint creation: {e}")

    # --- Step 2: Register a model (using an Azure OpenAI foundation model) ----
    console.print("\n[bold]Step 2:[/bold] Registering model...", style="cyan")

    model = Model(
        name="gpt-4o-mini-deployment",
        version="1",
        path="azureml://registries/azure-openai/models/gpt-4o-mini/versions/1",
        type="custom_model",
        description="GPT-4o-mini model for managed endpoint demo",
    )

    try:
        registered_model = ml_client.models.create_or_update(model)
        console.print(f"   ✅ Model registered: {registered_model.name} v{registered_model.version}")
    except Exception as e:
        console.print(f"   ⚠️  Model registration: {e}")
        registered_model = model

    # --- Step 3: Create a deployment with scaling configuration ---------------
    console.print("\n[bold]Step 3:[/bold] Creating deployment...", style="cyan")

    deployment = ManagedOnlineDeployment(
        name="blue",
        endpoint_name=endpoint_name,
        model=registered_model,
        instance_type="Standard_DS3_v2",
        instance_count=1,
        request_settings=OnlineRequestSettings(
            request_timeout_ms=90000,
            max_concurrent_requests_per_instance=10,
        ),
        liveness_probe=ProbeSettings(
            initial_delay=600,
            period=10,
            timeout=2,
            failure_threshold=30,
        ),
        readiness_probe=ProbeSettings(
            initial_delay=10,
            period=10,
            timeout=2,
            failure_threshold=30,
        ),
    )

    try:
        ml_client.online_deployments.begin_create_or_update(deployment).result()
        console.print(f"   ✅ Deployment 'blue' created with 1 instance (Standard_DS3_v2)")
    except Exception as e:
        console.print(f"   ⚠️  Deployment creation: {e}")

    # --- Step 4: Route 100% traffic to the deployment ------------------------
    console.print("\n[bold]Step 4:[/bold] Routing traffic...", style="cyan")

    try:
        endpoint.traffic = {"blue": 100}
        ml_client.online_endpoints.begin_create_or_update(endpoint).result()
        console.print("   ✅ 100% traffic routed to 'blue' deployment")
    except Exception as e:
        console.print(f"   ⚠️  Traffic routing: {e}")

    # --- Step 5: Get endpoint details ----------------------------------------
    console.print("\n[bold]Step 5:[/bold] Endpoint details:", style="cyan")

    try:
        ep = ml_client.online_endpoints.get(endpoint_name)
        console.print(f"   Scoring URI : {ep.scoring_uri}")
        console.print(f"   Swagger URI : {ep.openapi_uri}")
        console.print(f"   Provisioning: {ep.provisioning_state}")
        console.print(f"   Traffic     : {ep.traffic}")
    except Exception as e:
        console.print(f"   ⚠️  Could not fetch details: {e}")

    console.print("\n[bold green]✅ Deployment pipeline complete![/bold green]")
    console.print(
        "\n💡 To clean up, run:\n"
        f"   az ml online-endpoint delete --name {endpoint_name} --yes"
    )


if __name__ == "__main__":
    main()