from enum import Enum
from typing import Optional, Union

from pydantic import BaseModel, Field

from weave.trace_server.interface.builtin_object_classes import base_object_def


class ResponseFormat(str, Enum):
    JSON = "json"
    TEXT = "text"

    # TODO: Fast follow up
    # JSON_SCHEMA = "jsonschema"


class Message(BaseModel):
    """A message in a conversation with an LLM.

    Attributes:
        role: The role of the message's author. Can be: system, user, assistant, function or tool.
        content: The contents of the message. Required for all messages, but may be null for assistant messages with function calls.
        name: The name of the author of the message. Required if role is "function". Must match the name of the function represented in content.
              Can contain characters (a-z, A-Z, 0-9), and underscores, with a maximum length of 64 characters.
        function_call: The name and arguments of a function that should be called, as generated by the model.
        tool_call_id: Tool call that this message is responding to.
    """

    role: str
    content: Optional[Union[str, list[dict]]]
    name: Optional[str]
    function_call: Optional[dict]
    tool_call_id: Optional[str]


class LLMStructuredCompletionModelDefaultParams(BaseModel):
    # Could use Prompt objects for the message template
    # This is a list of Messages, loosely following litellm's message format
    # https://docs.litellm.ai/docs/completion/input#properties-of-messages
    messages_template: list[Message]

    temperature: Optional[float] = None
    top_p: Optional[float] = None
    max_tokens: Optional[int] = None
    presence_penalty: Optional[float] = None
    frequency_penalty: Optional[float] = None
    stop: Optional[list[str]] = None
    n_times: Optional[int] = None
    functions: Optional[list[dict]] = None

    # Either json, text, or json_schema
    response_format: Optional[ResponseFormat] = None

    # TODO: Currently not used. Fast follow up with json_schema
    # if default_params.response_format is set to JSON_SCHEMA, this will be used
    # response_format_schema: dict | None = None


class LLMStructuredCompletionModel(base_object_def.BaseObject):
    # <provider>/<model> or ref to a provider model
    llm_model_id: Union[str, base_object_def.RefStr]

    default_params: LLMStructuredCompletionModelDefaultParams = Field(
        default_factory=LLMStructuredCompletionModelDefaultParams
    )
