package main

import (
    "context"
    "fmt"
    "log"
    "os"
    "github.com/Azure/azure-sdk-for-go/sdk/ai/azopenai"
    "github.com/Azure/azure-sdk-for-go/sdk/azcore"
    "github.com/Azure/azure-sdk-for-go/sdk/azcore/to"
)

func main() {
    // Retrieve endpoint and API key from environment variables
    azureOpenAIEndpoint := os.Getenv("AZURE_OPENAI_ENDPOINT")
    azureOpenAIKey := os.Getenv("AZURE_OPENAI_API_KEY")
    deploymentName := "o4-mini"

    // Define the parameters
    maxTokens := int32(100000)

    // Check if the endpoint and API key environment variables are set
    if azureOpenAIEndpoint == "" || azureOpenAIKey == "" {
        fmt.Fprintf(os.Stderr, "Please set the AZURE_OPENAI_ENDPOINT and AZURE_OPENAI_API_KEY environment variables.\n")
        return
    }

    // Initialize the OpenAI client with API key-based authentication
    cred := azcore.NewKeyCredential(azureOpenAIKey)
    client, err := azopenai.NewClientWithKeyCredential(azureOpenAIEndpoint, cred, nil)
    if err != nil {
        log.Fatalf("ERROR: %s", err)
    }

    // Define the messages for the chat interaction
    messages := []azopenai.ChatRequestMessageClassification{
        &azopenai.ChatRequestDeveloperMessage{Content: azopenai.NewChatRequestDeveloperMessageContent("You are an AI assistant that helps people find information.")},
    }

    // Make the chat completion request
    resp, err := client.GetChatCompletions(context.TODO(), azopenai.ChatCompletionsOptions{
        Messages:       messages,
        DeploymentName: &deploymentName,
        MaxCompletionTokens:      &maxTokens,
    }, nil)
    if err != nil {
        log.Fatalf("ERROR: %s", err)
    }

    // Print the response
    for _, choice := range resp.Choices {
        if choice.Message != nil && choice.Message.Content != nil {
            fmt.Printf("Response[%d]: %s\n", *choice.Index, *choice.Message.Content)
        }
        if choice.FinishReason != nil {
            fmt.Printf("Finish reason[%d]: %s\n", *choice.Index, *choice.FinishReason)
        }
    }
}