"""
deploy_model.py — Deploy a model and run a chat completion
==========================================================
Demonstrates how to deploy a model through the Azure AI Foundry SDK and then
call it for a chat completion.

Prerequisites:
    pip install azure-ai-projects azure-ai-inference azure-identity python-dotenv

Usage:
    python deploy_model.py
"""

import os
import sys
from dotenv import load_dotenv
from azure.identity import DefaultAzureCredential
from azure.ai.projects import AIProjectClient
from azure.ai.inference import ChatCompletionsClient
from azure.ai.inference.models import SystemMessage, UserMessage

load_dotenv()


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


def main() -> None:
    """Connect to a project, obtain a chat client, and run a completion."""

    credential = DefaultAzureCredential()

    # Build the project connection string
    location = get_env("AZURE_LOCATION")
    conn_str = (
        f"{location}.api.azureml.ms;"
        f"{get_env('AZURE_SUBSCRIPTION_ID')};"
        f"{get_env('AZURE_RESOURCE_GROUP')};"
        f"{get_env('AZURE_AI_FOUNDRY_PROJECT_NAME')}"
    )

    project_client = AIProjectClient.from_connection_string(
        credential=credential,
        conn_str=conn_str,
    )

    # Get a ChatCompletionsClient from the project's default connection
    deployment_name = get_env("AZURE_OPENAI_CHAT_DEPLOYMENT")
    print(f"Using deployment: {deployment_name}")

    chat_client: ChatCompletionsClient = (
        project_client.inference.get_chat_completions_client()
    )

    # Run a simple chat completion
    response = chat_client.complete(
        model=deployment_name,
        messages=[
            SystemMessage(content="You are a helpful AI assistant for Azure learners."),
            UserMessage(content="What is Azure AI Foundry and why should I use it?"),
        ],
        temperature=0.7,
        max_tokens=512,
    )

    print("\n🤖 Model response:")
    print(response.choices[0].message.content)
    print(f"\nTokens used — prompt: {response.usage.prompt_tokens}, "
          f"completion: {response.usage.completion_tokens}")


if __name__ == "__main__":
    main()