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


async def llm_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 = 600,
    web_search = False
) -> Union[str, BaseModel]:
    """
    Async unified interface to call different LLM APIs with optional structured output.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    content : str
        User-provided prompt or content for which a response is requested.
    instructions : str
        System-level instructions to guide the model's behavior and response generation.
    model_name : str, optional
        The identifier for the target model to be used.
        Defaults to "gpt-4o".
    response_format : Optional[Type[BaseModel]], optional
        Pydantic model class to enforce structured JSON output.
        If None, returns regular string response.

    ---------------------------------------------------------------------------
    OUTPUT:
    output : Union[str, BaseModel]
        The response - either a string (regular) or Pydantic model instance (structured).
    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    available_models = bs.read_data("api_model_dict.json")
    #available_open_router_models = bs.read_data('openrouter_models.json')
    
    if (
        (model_name in available_models["openai"])
        or model_name in ("deepseek")
        or model_name in available_models["deepseek"]
    ):
        output = await bs.openai_api_async(
            content=content,
            instructions=instructions,
            model_name=model_name,
            response_format=response_format,
            temperature=temperature,
            max_tokens=max_tokens,
            mymessages=mymessages,
            timeout = timeout,
            web_search = web_search
        )

    elif model_name in available_models["claude"] or model_name in ("claude"):
        if not max_tokens:
            max_tokens = 1000
        output = await bs.claude_api_async(
            content=content,
            instructions=instructions,
            model_name=model_name,
            response_format=response_format,
            temperature=temperature,
            max_tokens=max_tokens
        )

    elif model_name in available_models["gemini"] or model_name in ("gemini"):
        output = await bs.gemini_api_async(
            content=content,
            instructions=instructions,
            model_name=model_name,
            response_format=response_format,
            temperature=temperature,
            max_tokens=max_tokens
        )
    elif model_name[:7] == 'ollama/':
            output = await bs.ollama_async(
                content=content,
                instructions=instructions,
                model_name=model_name,
            )
    else:
        output = await bs.openai_api_async(content, instructions, model_name, 
                                        temperature=temperature, max_tokens=max_tokens, mymessages=mymessages, timeout = timeout, response_format=response_format, web_search = web_search)
    
    return output