"""
chat_completions.py — Comprehensive Chat Completions Examples
=============================================================
Covers four patterns using the Azure OpenAI Python SDK:
  1. Standard (non-streaming) completion
  2. Streaming completion
  3. Function calling (tool use)
  4. JSON mode (structured output)

Prerequisites:
    pip install openai azure-identity python-dotenv

Usage:
    python chat_completions.py
"""

import json
import os
import sys
from dotenv import load_dotenv
from openai import AzureOpenAI
from azure.identity import DefaultAzureCredential, get_bearer_token_provider

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 get_client() -> AzureOpenAI:
    """Create an AzureOpenAI client using Entra ID (recommended) or API key."""
    endpoint = get_env("AZURE_OPENAI_ENDPOINT")
    api_version = get_env("AZURE_OPENAI_API_VERSION")
    api_key = os.getenv("AZURE_OPENAI_API_KEY")

    if api_key:
        return AzureOpenAI(
            azure_endpoint=endpoint,
            api_key=api_key,
            api_version=api_version,
        )

    # Token-based auth via DefaultAzureCredential
    token_provider = get_bearer_token_provider(
        DefaultAzureCredential(),
        "https://cognitiveservices.azure.com/.default",
    )
    return AzureOpenAI(
        azure_endpoint=endpoint,
        azure_ad_token_provider=token_provider,
        api_version=api_version,
    )


# ---------------------------------------------------------------------------
# 1. Standard (non-streaming) completion
# ---------------------------------------------------------------------------
def demo_standard(client: AzureOpenAI, deployment: str) -> None:
    print("\n" + "=" * 60)
    print("1️⃣  Standard Chat Completion")
    print("=" * 60)

    response = client.chat.completions.create(
        model=deployment,
        messages=[
            {"role": "system", "content": "You are a concise Azure expert."},
            {"role": "user", "content": "Explain Azure AI Foundry in 3 sentences."},
        ],
        temperature=0.7,
        max_tokens=256,
    )
    print(response.choices[0].message.content)
    print(f"(tokens — prompt: {response.usage.prompt_tokens}, "
          f"completion: {response.usage.completion_tokens})")


# ---------------------------------------------------------------------------
# 2. Streaming completion
# ---------------------------------------------------------------------------
def demo_streaming(client: AzureOpenAI, deployment: str) -> None:
    print("\n" + "=" * 60)
    print("2️⃣  Streaming Chat Completion")
    print("=" * 60)

    stream = client.chat.completions.create(
        model=deployment,
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": "List 5 benefits of using Azure AI Foundry."},
        ],
        temperature=0.7,
        max_tokens=512,
        stream=True,
    )

    for chunk in stream:
        if chunk.choices and chunk.choices[0].delta.content:
            print(chunk.choices[0].delta.content, end="", flush=True)
    print()  # newline after stream


# ---------------------------------------------------------------------------
# 3. Function calling (tool use)
# ---------------------------------------------------------------------------
def get_weather(location: str, unit: str = "celsius") -> dict:
    """Simulated weather lookup for the function-calling demo."""
    return {"location": location, "temperature": 22, "unit": unit, "condition": "Sunny"}


WEATHER_TOOL = {
    "type": "function",
    "function": {
        "name": "get_weather",
        "description": "Get the current weather for a location.",
        "parameters": {
            "type": "object",
            "properties": {
                "location": {"type": "string", "description": "City name, e.g. 'Seattle'"},
                "unit": {"type": "string", "enum": ["celsius", "fahrenheit"], "default": "celsius"},
            },
            "required": ["location"],
        },
    },
}


def demo_function_calling(client: AzureOpenAI, deployment: str) -> None:
    print("\n" + "=" * 60)
    print("3️⃣  Function Calling (Tool Use)")
    print("=" * 60)

    messages = [
        {"role": "system", "content": "You are a helpful assistant that can look up weather."},
        {"role": "user", "content": "What's the weather in Seattle?"},
    ]

    # First call — the model may ask to call our tool
    response = client.chat.completions.create(
        model=deployment,
        messages=messages,
        tools=[WEATHER_TOOL],
        tool_choice="auto",
    )

    assistant_msg = response.choices[0].message

    if assistant_msg.tool_calls:
        # Execute each tool call
        messages.append(assistant_msg)
        for tool_call in assistant_msg.tool_calls:
            args = json.loads(tool_call.function.arguments)
            result = get_weather(**args)
            messages.append({
                "role": "tool",
                "tool_call_id": tool_call.id,
                "content": json.dumps(result),
            })

        # Second call — model generates a natural-language answer
        follow_up = client.chat.completions.create(
            model=deployment,
            messages=messages,
            tools=[WEATHER_TOOL],
        )
        print(follow_up.choices[0].message.content)
    else:
        print(assistant_msg.content)


# ---------------------------------------------------------------------------
# 4. JSON mode (structured output)
# ---------------------------------------------------------------------------
def demo_json_mode(client: AzureOpenAI, deployment: str) -> None:
    print("\n" + "=" * 60)
    print("4️⃣  JSON Mode (Structured Output)")
    print("=" * 60)

    response = client.chat.completions.create(
        model=deployment,
        messages=[
            {
                "role": "system",
                "content": (
                    "You are a helpful assistant that responds in JSON. "
                    "Return an object with keys: services (array of strings), "
                    "summary (string)."
                ),
            },
            {
                "role": "user",
                "content": "List the top 3 Azure AI services and summarize them.",
            },
        ],
        response_format={"type": "json_object"},
        temperature=0.3,
        max_tokens=512,
    )

    raw = response.choices[0].message.content
    parsed = json.loads(raw)
    print(json.dumps(parsed, indent=2))


# ---------------------------------------------------------------------------
# Main — run all demos
# ---------------------------------------------------------------------------
def main() -> None:
    client = get_client()
    deployment = get_env("AZURE_OPENAI_CHAT_DEPLOYMENT")

    demo_standard(client, deployment)
    demo_streaming(client, deployment)
    demo_function_calling(client, deployment)
    demo_json_mode(client, deployment)

    print("\n✅ All demos completed!")


if __name__ == "__main__":
    main()