from typing import Optional, Type, Union
import pydantic
from pydantic import BaseModel
import bilbystats as bs


async def openai_api_async(
    content: str,
    instructions: str = "",
    model_name: str = "gpt-4o",
    response_format: Optional[Type[BaseModel]] = None,
    temperature: float = 1.0,
    max_tokens: int = None,
    mymessages = None,
    timeout = 5,
    web_search = False
) -> Union[str, BaseModel]:
    
    if model_name == "deepseek":
        model_name = "deepseek-chat"


    # Get reusable client instead of creating new one
    client = await bs.get_openai_client(model_name, timeout = timeout)
    
    if model_name.endswith(':groq'):
        model_name = model_name[:-5] 
    
    available_models = bs.read_data("api_model_dict.json")

    messages = [] 
    if len(instructions) > 0:
        messages.append({"role": "system", "content": instructions})
        
    if mymessages:
        messages = messages + mymessages
        
    messages.append({"role": "user","content": content})

    completion_args = {
        "model": model_name,
        "messages": messages,
        "temperature": temperature,
        "max_tokens": max_tokens,
        "timeout": timeout 
    }

    if response_format:
       completion_args["response_format"] = response_format
        
    if web_search:
       completion_args["tools"] = [{"type": "web_search"}]
    
    if type(response_format) == pydantic._internal._model_construction.ModelMetaclass:
        if model_name in available_models["deepseek"]:
            response_format = {
                "type": "json_object",
            }
        completion_args["response_format"] = response_format

        completion = await client.beta.chat.completions.parse(**completion_args)

        if model_name in available_models["openai"]:
            output = completion.choices[0].message.parsed
        elif model_name in available_models["deepseek"]:
            output = completion.choices[0].message.content
    else:
        completion = await client.chat.completions.create(**completion_args)
        output = completion.choices[0].message.content

    return output