from enum import Enum
from typing import Optional
from pydantic import BaseModel, Field


class OpenAIModel(str, Enum):
    GPT_4 = "gpt-4"
    GPT_4_0613 = "gpt-4-0613"
    GPT_4_32K = "gpt-4-32k"
    GPT_4_32K_0613 = "gpt-4-32k-0613"
    GPT_4_1106_PREVIEW= "gpt-4-1106-preview"
    GPT_3_5_TURBO = "gpt-3.5-turbo"
    GPT_3_5_TURBO_16K = "gpt-3.5-turbo-16k"
    GPT_3_5_TURBO_INSTRUCT = "gpt-3.5-turbo-instruct"
    GPT_3_5_TURBO_0613 = "gpt-3.5-turbo-0613"
    GPT_3_5_TURBO_16K_0613 = "gpt-3.5-turbo-16k-0613"
    TEXT_MODERATION_LATEST = "text-moderation-latest"
    TEXT_MODERATION_STABLE = "text-moderation-stable"


class Message(BaseModel):
    """
    Message for OpenAI chat.
    """

    class FunctionCall(BaseModel):
        """
        The name and arguments of a function that should be called, as generated by the model.
        """

        name: str = Field(description="Name of the function to call.")
        arguments: str = Field(
            description="Arguments to call the function with, as generated by the model in JSON format."
        )

    role: Optional[str] = Field(description="Role of the message.")
    name: Optional[str] = Field(description="Name of the message.")
    content: Optional[str] = Field(description="Content of the message.")
    function_call: Optional[FunctionCall] = Field(
        description="Function call of the message."
    )


class FunctionDefinition(BaseModel):
    """
    Function definition for OpenAI function call.
    """

    class FunctionParameter(BaseModel):
        """
        Parameter for OpenAI function call.
        """

        name: str = Field(description="Name of the parameter.")
        description: str = Field(description="Description of the parameter.")
        required: bool = Field(
            default=True, description="Whether the parameter is required or not."
        )

    name: str = Field("Name of the function.")
    description: str = Field("Description of the function.")
    parameters: list[FunctionParameter] = Field("Parameters of the function.")

    def dict(self) -> dict[str, any]:
        return {
            "name": self.name,
            "description": self.description,
            "parameters": {
                "type": "object",
                "properties": {
                    param.name: {
                        "type": "string",
                        "description": param.description,
                    }
                    for param in self.parameters
                },
                "required": [param.name for param in self.parameters if param.required],
            },
        }


class ChatConfig(BaseModel):
    model: OpenAIModel = Field(description="Model to use for chat.")
    use_streaming: bool = Field(
        default=False, description="Whether to use streaming output or not."
    )
    append_history: bool = Field(
        default=False, description="Whether to append this message to history."
    )

class OldCompleteConfig(BaseModel):
    model: OpenAIModel = Field(description="Model to use for chat.")
    use_streaming: bool = Field(
        default=False, description="Whether to use streaming output or not."
    )
    

class OldCompleteInput(OldCompleteConfig):
    prompt:str = Field(description="prompt to use for model.")

class CompleteInput(ChatConfig):
    pass


class ChatInput(ChatConfig):
    message_text: str = Field(description="Message to chat.")


class ChatWithPromptTemplateInput(ChatConfig):
    prompt_template: str = Field(description="Prompt template for chat.")
    params: dict[str, str] = Field(description="Parameters for prompt template.")


class ChatWithMessageInput(ChatConfig):
    message: Message = Field(description="Message to chat.")


class ChatWithMessagesInput(ChatConfig):
    messages: list[Message] = Field(description="Messages to chat.")


class OpenAIResp(BaseModel):
    """
    Response from OpenAI with no streaming. It's a chat completion object.
    """

    message: Message = Field(description="Response message")
    finish_reason: str = Field(description="Finish reason.")


# TODO: support streaming response
class OpenAIStreamingResp(BaseModel):
    """
    Response from OpenAI with streaming.
    """

    delta: Message = Field(description="Delta message")
    finish_reason: str = Field(description="Finish reason.")

class OpenAIOldResp(BaseModel):
    """
    Response from OpenAI with no streaming. It's a chat completion object.
    """
    text: str = Field(description="Response message")
    finish_reason: str = Field(description="Finish reason.")
