from typing import Any
from dialtone.types import ChatMessage, Tool, DialtoneClient


def prepare_chat_message(message: ChatMessage | dict[str, Any]) -> dict:
    if isinstance(message, dict):
        return message

    model_dump = message.model_dump()

    # Have to manually remove this field because an empty array
    # here isn't the same as an empty array generated by the backend
    if "tool_calls" in model_dump and len(model_dump["tool_calls"]) == 0:
        del model_dump["tool_calls"]
    # Same for name field
    if "name" in model_dump and not model_dump["name"]:
        del model_dump["name"]

    return model_dump


def prepare_tool(tool: Tool | dict[str, Any]) -> dict:
    if isinstance(tool, dict):
        return tool

    return tool.model_dump()


def prepare_chat_completion(
    client: DialtoneClient,
    messages: list[ChatMessage] | list[dict[str, Any]],
    stream: bool = False,
    tools: list[Tool] | list[dict] = [],
) -> tuple[dict, dict]:
    headers = {"Authorization": f"Bearer {client.api_key}"}
    params = {
        "messages": [prepare_chat_message(message) for message in messages],
        "dials": client.dials.model_dump(),
        "provider_config": client.provider_config.model_dump(),
    }
    if stream:
        params["stream"] = True
    if client.router_model_config:
        params["router_model_config"] = client.router_model_config.model_dump()
    if client.fallback_config:
        params["fallback_config"] = client.fallback_config.model_dump()
    if client.tools_config:
        params["tools_config"] = client.tools_config.model_dump()
    if tools:
        params["tools"] = [prepare_tool(tool) for tool in tools]

    return headers, params


def prepare_chat_route(
    client: DialtoneClient,
    messages: list[ChatMessage] | list[dict[str, Any]],
    tools: list[Tool] = [],
) -> tuple[dict, dict]:
    headers = {"Authorization": f"Bearer {client.api_key}"}
    params = {
        "messages": [prepare_chat_message(message) for message in messages],
        "dials": client.dials.model_dump(),
        "provider_config": client.provider_config.model_dump(),
    }
    if client.router_model_config:
        params["router_model_config"] = client.router_model_config.model_dump()
    if client.fallback_config:
        params["fallback_config"] = client.fallback_config.model_dump()
    if client.tools_config:
        params["tools_config"] = client.tools_config.model_dump()
    if tools:
        params["tools"] = [tool.model_dump() for tool in tools]

    return headers, params
