from typing import List, Type, Optional
from pydantic import BaseModel, Field
import json
import os
from azure.ai.inference import ChatCompletionsClient
from azure.ai.inference.models import SystemMessage, UserMessage
from azure.core.credentials import AzureKeyCredential

endpoint = "https://d-robotics.openai.azure.com/openai/deployments/gpt-4o"
model_name = "gpt-4o"

# Get API key from environment variable
api_key = os.environ.get("AZURE_API_KEY")
if not api_key:
    raise ValueError("AZURE_API_KEY environment variable is required but not set")

client = ChatCompletionsClient(
    endpoint=endpoint,
    credential=AzureKeyCredential(api_key),
)


def generate(messages: List[dict], custom_format: Type[BaseModel]) -> Optional[BaseModel]:
    strformat = custom_format.schema_json()
    messages.append({
        "role": "system",
        "content": "you shall output a json object with the following format: " + strformat,
    })
    response = client.complete(
        messages=messages,
        max_tokens=4096,
        temperature=0.8,
        top_p=1.0,
        model=model_name,
        response_format="json_object",
    )

    json_content = response.choices[0].message.content
    if json_content:
        parsed_json = json.loads(json_content)
        return (custom_format.parse_obj(parsed_json)
                if hasattr(custom_format, "parse_obj") else custom_format.model_validate(parsed_json))

    return None


if __name__ == "__main__":
    pass
