from typing import Any, Dict, List, Literal, Optional, Union

from pydantic import BaseModel, field_serializer

# =============== OpenAI ===============
# We have ChatCompletion and Response APIs


# =============== ChatCompletion API ===============
class Function(BaseModel):
    arguments: str
    """
    The arguments to call the function with, as generated by the model in JSON
    format. Note that the model does not always generate valid JSON, and may
    hallucinate parameters not defined by your function schema. Validate the
    arguments in your code before calling your function.
    """

    name: str
    """The name of the function to call."""


class ChatCompletionMessageToolCall(BaseModel):
    """
    Example usage

    {
        "id": "call_12345xyz",
        "type": "function",
        "function": {
            "name": "get_weather",
            "arguments": "{\"location\":\"Paris, France\"}"
        }
    }
    """

    id: str
    """The ID of the tool call."""

    function: Function
    """The function that the model called."""

    type: Literal["function"] = "function"
    """The type of the tool. Currently, only `function` is supported."""


class ChatCompetionMessageToolCallResult(BaseModel):
    """
    Example usage

    {
        "role": "tool",
        "tool_call_id": tool_call.id,
        "content": str(result)
    }
    """

    role: Literal["tool"] = "tool"
    """The role of the message. Always `tool`."""

    tool_call_id: str
    """The ID of the tool call that this message is responding to."""

    content: str
    """Result of the tool call in string format."""


class ChatCompletionMessage(BaseModel):
    content: Optional[str] = None
    """The contents of the message."""

    role: Literal["assistant"] = "assistant"
    """The role of the author of this message."""

    tool_calls: Optional[List[ChatCompletionMessageToolCall]] = None
    """The tool calls generated by the model, such as function calls."""


# =============== Reponse API ===============
class ResponseFunctionToolCall(BaseModel):
    """
    Example usage

    {
        "id": "fc_12345xyz",
        "call_id": "call_12345xyz",
        "type": "function_call",
        "name": "get_weather",
        "arguments": "{\"location\":\"Paris, France\"}"
    }
    """

    arguments: str
    """A JSON string of the arguments to pass to the function."""

    call_id: str
    """The unique ID of the function tool call generated by the model."""

    name: str
    """The name of the function to run."""

    type: Literal["function_call"] = "function_call"
    """The type of the function tool call. Always `function_call`."""

    id: Optional[str] = None
    """The unique ID of the function tool call."""

    status: Optional[Literal["in_progress", "completed", "incomplete"]] = None
    """The status of the item.

    One of `in_progress`, `completed`, or `incomplete`. Populated when items are
    returned via API.
    """


class ResponseFunctionToolCallResult(BaseModel):
    """
    Example usage

    {
        "type": "function_call_output",
        "call_id": tool_call.call_id,
        "output": str(result)
    }
    """

    type: Literal["function_call_output"] = "function_call_output"
    """The type of the function tool call result. Always `function_call_output`."""

    call_id: str
    """The unique ID of the function tool call."""

    output: str
    """The output of the function tool call as a string."""


# =============== Common ===============
class ToolCall(BaseModel):
    id: str
    """The ID of the tool call."""
    name: str
    """The name of the function to call."""
    arguments: str
    """The arguments to call the function with, as generated by the model in JSON format."""

    @classmethod
    def from_tool_call(cls, tool_call: Any) -> "ToolCall":
        if resemble_type(tool_call, ChatCompletionMessageToolCall):
            return cls(
                id=tool_call.id,
                name=tool_call.function.name,
                arguments=tool_call.function.arguments,
            )
        elif resemble_type(tool_call, ResponseFunctionToolCall):
            return cls(
                id=tool_call.call_id,
                name=tool_call.name,
                arguments=tool_call.arguments,
            )
        else:
            raise TypeError("Unsupported type for conversion")


class ToolCallResult(BaseModel):
    id: str
    """The ID of the tool call."""
    result: Any
    """The result of the tool call."""

    @field_serializer("result")
    def convert_any_field_to_str(self, value: Any) -> str:
        """Convert `result` to string during serialization.

        Args:
            value: Current value of `any_field`.

        Returns:
            str: String representation of the field.
        """
        # simple case for now, can be extended to handle more complex cases
        return str(value)


API_FORMATS = Literal[
    "openai",  # old default, alias to openai-chatcompletion
    "openai-chatcompletion",  # chat completion
    "openai-response",
    "anthropic",
    "gemini",
]


def resemble_type(obj: object, cls: type) -> bool:
    """Check if the object's class matches the given class type.

    Args:
        obj (object): The object to check.
        cls (type): The class type to compare against.

    Returns:
        bool: True if the object's class name matches the given class type name, False otherwise.
    """
    class_name = obj.__class__.__name__

    if class_name == cls.__name__:
        return True
    return False


def convert_tool_calls(tool_calls: List[Any]) -> List[ToolCall]:
    """Convert a list of tool calls into a list of ToolCall objects.

    Args:
        tool_calls (List[Any]): A list of tool call objects to convert.

    Returns:
        List[ToolCall]: A list of converted ToolCall objects.
    """
    return [ToolCall.from_tool_call(tool_call) for tool_call in tool_calls]


def recover_assistant_message(
    tool_calls: List[ToolCall],
    *,
    api_format: API_FORMATS = "openai",
) -> List[Dict[str, Any]]:
    """Recover the assistant message from tool calls in various API formats.

    Args:
        tool_calls (List[ToolCall]): A list of ToolCall objects.
        api_format (API_FORMATS, optional): The desired API format. Defaults to "openai".

    Returns:
        List[Dict[str, Any]]: The assistant message in the specified API format.

    Raises:
        NotImplementedError: If the API format is "anthropic" or "gemini".
        ValueError: If the API format is unsupported.
    """
    if api_format in ["openai", "openai-chatcompletion"]:
        message = [
            ChatCompletionMessage(
                tool_calls=[
                    ChatCompletionMessageToolCall(
                        id=tool_call.id,
                        function=Function(
                            name=tool_call.name,
                            arguments=tool_call.arguments,
                        ),
                    )
                    for tool_call in tool_calls
                    if tool_call.name and tool_call.arguments
                ]
            ).model_dump()
        ]

        return message

    elif api_format == "openai-response":
        message = [
            ResponseFunctionToolCall(
                call_id=tool_call.id,
                name=tool_call.name,
                arguments=tool_call.arguments,
            ).model_dump()
            for tool_call in tool_calls
        ]

        return message

    elif api_format == "anthropic":
        raise NotImplementedError("Anthropic API format is not supported yet.")
    elif api_format == "gemini":
        raise NotImplementedError("Gemini API format is not supported yet.")
    else:
        raise ValueError(f"Unsupported API format: {api_format}")


def recover_tool_message(
    tool_responses: Dict[str, str],
    *,
    api_format: API_FORMATS = "openai",
) -> List[Dict[str, Any]]:
    """Recover the tool message from tool responses in various API formats.

    Args:
        tool_responses (Dict[str, str]): A dictionary of tool responses with call_ids as keys and results as values.
        api_format (API_FORMATS, optional): The desired API format. Defaults to "openai".

    Returns:
        List[Dict[str, Any]]: A list of tool messages in the specified API format.

    Raises:
        NotImplementedError: If the API format is "anthropic" or "gemini".
        ValueError: If the API format is unsupported.
    """
    messages = []
    for call_id, result in tool_responses.items():
        if api_format in ["openai", "openai-chatcompletion"]:
            tool_message = ChatCompetionMessageToolCallResult(
                tool_call_id=call_id,
                content=str(result),
            )
            messages.append(tool_message.model_dump())
        elif api_format == "openai-response":
            tool_message = ResponseFunctionToolCallResult(
                call_id=call_id,
                output=str(result),
            )
            messages.append(tool_message.model_dump())
        elif api_format == "anthropic":
            raise NotImplementedError("Anthropic API format is not supported yet")
        elif api_format == "gemini":
            raise NotImplementedError("Gemini API format is not supported yet")
        else:
            raise ValueError(f"Unsupported API format: {api_format}")

    return messages
