from typing import Protocol, cast

from openai import AsyncAzureOpenAI
from openai.types.chat import ChatCompletion


class AsyncGPT:
    def __init__(self, azure_endpoint: str, api_key: str, api_version: str) -> None:
        self.client = AsyncAzureOpenAI(
            azure_endpoint=azure_endpoint, api_key=api_key, api_version=api_version
        )

    async def generate(self, messages, model, **kwargs):
        response = await self.client.chat.completions.create(
            messages=messages, model=model, **kwargs
        )

        response = cast(ChatCompletion, response)

        return response.choices[0].message

    async def stream(self, messages, model, **kwargs):
        response = await self.client.chat.completions.create(
            messages=messages, model=model, stream=True, **kwargs
        )

        async for chunk in response:
            if choices := chunk.choices:
                yield choices[0].delta

    @staticmethod
    def build_tool_call(tool_id: str, tool_name: str, tool_buffer: str):
        tool_call = {
            "tool_calls": [
                {
                    "id": tool_id,
                    "function": {
                        "name": tool_name,
                        "arguments": tool_buffer,
                    },
                    "type": "function",
                }
            ]
        }

        return tool_call

    @staticmethod
    def build_tool_call_id(tool_id: str):
        return {"tool_call_id": tool_id}

    class Config(Protocol):
        azure_endpoint: str
        openai_api_key: str
        openai_api_version: str

    @classmethod
    def from_config(cls, config: Config):
        return cls(
            config.azure_endpoint, config.openai_api_key, config.openai_api_version
        )
