import json
import time
import typing
from openai import AsyncOpenAI
from openai.types.chat import ChatCompletionToolParam
from .base import DEFAULT_MODEL_CONFIGURATION, Conversation, ModelConfiguration
from pydantic import BaseModel


class Tool(BaseModel):
    name: str
    description: str
    parameters: type[BaseModel]


class ToolCall(BaseModel):
    id: str
    name: str
    arguments: dict


class ToolCallCompletion(BaseModel):
    conversation: Conversation
    start_time: float
    end_time: float
    tool_calls: list[ToolCall]


def tool_to_openai_tool(tool: Tool) -> ChatCompletionToolParam:
    model_json_schema = tool.parameters.model_json_schema()
    model_json_schema["additionalProperties"] = False

    return typing.cast(
        ChatCompletionToolParam,
        {
            "type": "function",
            "function": {
                "name": tool.name,
                "description": tool.description,
                "parameters": model_json_schema,
            },
            "strict": True,
        },
    )


async def generate_tool_call_completion_for_conversation(
    conversation: Conversation,
    tools: Tool | list[Tool],
    model_configuration: ModelConfiguration | None = None,
) -> ToolCallCompletion:
    """Generate a tool call completion from a conversation.

    This function sends the conversation to the OpenAI API with the provided tools,
    and returns the tool calls that the model generates.

    Args:
        conversation: The conversation to generate a tool call completion for.
        tools: A Tool object or a list of Tool objects that define the available tools for the model.
        model_configuration: Optional configuration for the model. If None, uses the default configuration.

    Returns:
        A ToolCallCompletion object containing the tool calls generated by the model.
    """
    if model_configuration is None:
        model_configuration = DEFAULT_MODEL_CONFIGURATION

    if isinstance(tools, Tool):
        tools = [tools]

    client = AsyncOpenAI(
        api_key=model_configuration.api_key,
        base_url=model_configuration.base_url,
    )

    openai_tools = [tool_to_openai_tool(tool) for tool in tools]

    start_time = time.time()
    completion = await client.chat.completions.create(
        model=model_configuration.model,
        messages=conversation.openai_messages,
        temperature=model_configuration.temperature,
        logprobs=model_configuration.log_probs,
        top_logprobs=model_configuration.top_log_probs,
        max_tokens=model_configuration.max_tokens,
        presence_penalty=model_configuration.presence_penalty,
        frequency_penalty=model_configuration.frequency_penalty,
        top_p=model_configuration.top_p,
        seed=model_configuration.seed,
        tools=openai_tools,
    )
    end_time = time.time()

    openai_tool_calls = completion.choices[0].message.tool_calls
    if openai_tool_calls is None:
        raise ValueError("No tool calls found in the completion.")

    tool_calls = [
        ToolCall(
            id=tool_call.id,
            name=tool_call.function.name,
            arguments=json.loads(tool_call.function.arguments),
        )
        for tool_call in openai_tool_calls
    ]

    return ToolCallCompletion(
        conversation=conversation,
        start_time=start_time,
        end_time=end_time,
        tool_calls=tool_calls,
    )


async def generate_tool_call_completion(
    user_prompt: str,
    tools: Tool | list[Tool],
    system_prompt: str | None = None,
    model_configuration: ModelConfiguration | None = None,
) -> ToolCallCompletion:
    """
    Generate a tool call completion for a user prompt.

    This function creates a conversation from the provided prompts and generates
    tool calls based on the available tools.

    Args:
        user_prompt: The user prompt to generate a tool call for.
        tools: A Tool object or a list of Tool objects that define the available tools for the model.
        system_prompt: Optional system prompt to use for the conversation.
        model_configuration: Optional configuration for the model. If None, uses the default configuration.

    Returns:
        A ToolCallCompletion object containing the generated tool calls.
    """
    if model_configuration is None:
        model_configuration = DEFAULT_MODEL_CONFIGURATION

    conversation = Conversation.from_prompts(user_prompt, system_prompt=system_prompt)

    return await generate_tool_call_completion_for_conversation(
        conversation,
        tools,
        model_configuration=model_configuration,
    )
