import enum
import uuid
from typing import Optional, List
from sqlalchemy import String, ForeignKey, JSON, Float, Table, DateTime, Column
from sqlalchemy.orm import Mapped, mapped_column, relationship
from datetime import datetime
from sqlalchemy.orm import declarative_base

Base = declarative_base()


def is_empty(value):
    if value is None:
        return True
    if isinstance(value, str):
        return value.strip() == ""
    if isinstance(value, list) or isinstance(value, dict):
        return len(value) == 0
    return False


def remove_empty_values(d):
    return {k: v for k, v in d.items() if not is_empty(v)}


class TimestampedModel(Base):
    """
    Base model class for all database models.
    """

    __abstract__ = True

    id: Mapped[str] = mapped_column(primary_key=True, default=lambda: str(uuid.uuid4()))
    namespace: Mapped[str] = mapped_column(String, default="default")
    created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.now)


class Model(TimestampedModel):
    """
    Defines a model used for inference.
    """

    __tablename__ = "models"

    prompts: Mapped[List["Prompts"]] = relationship("Prompt", back_populates="model")
    name: Mapped[str] = mapped_column(String)
    version: Mapped[Optional[str]] = mapped_column(String)
    provider: Mapped[Optional[str]] = mapped_column(String, nullable=True)
    description: Mapped[Optional[str]] = mapped_column(String, nullable=True)
    created: Mapped[datetime] = mapped_column(DateTime, default=datetime.now)

    def to_dict(self):
        """
        Convert the model to a dictionary format.
        """
        return remove_empty_values(
            {
                "id": self.id,
                "namespace": self.namespace,
                "name": self.name,
                "version": self.version,
                "provider": self.provider,
                "description": self.description,
                "created": self.created.timestamp(),
            }
        )


class MessageRole(enum.Enum):
    """
    Standard roles for messages in a conversation.
    """

    SYSTEM = "system"
    USER = "user"
    ASSISTANT = "assistant"
    TOOL = "tool"


class Message(TimestampedModel):
    """
    Defines a single message in a conversation.
    """

    __tablename__ = "chat_messages"

    prompt_id: Mapped[str] = mapped_column(ForeignKey("prompts.id"))
    prompt: Mapped["Prompt"] = relationship("Prompt", back_populates="messages")
    role: Mapped[MessageRole] = mapped_column(String)
    content: Mapped[Optional[str]] = mapped_column(String, nullable=True)
    tool_call_id: Mapped[Optional[str]] = mapped_column(String, nullable=True)

    def to_dict(self):
        """
        Convert the message to a dictionary format.
        """

        return remove_empty_values(
            {
                "role": self.role,
                "content": self.content,
                "tool_call_id": self.tool_call_id,
            }
        )


class ChatCompletion(TimestampedModel):
    """
    Defines a chat completion generated by a model.
    """

    __tablename__ = "chat_completions"

    prompt_id: Mapped[str] = mapped_column(ForeignKey("prompts.id"))
    prompt: Mapped["Prompt"] = relationship("Prompt", back_populates="choices")
    role: Mapped[MessageRole] = mapped_column(String)
    content: Mapped[Optional[str]] = mapped_column(String, nullable=True)
    tool_calls: Mapped[Optional[List["ToolCall"]]] = relationship(
        "ToolCall", back_populates="completion"
    )
    finish_reason: Mapped[Optional[str]] = mapped_column(String, nullable=True)

    def to_dict(self):
        """
        Convert the chat completion to a dictionary format.
        """

        return remove_empty_values(
            {
                "role": self.role,
                "content": self.content,
                "tool_calls": [tc.to_dict() for tc in self.tool_calls],
                "finish_reason": self.finish_reason,
            }
        )


prompt_tool_table = Table(
    "prompt_tool_association",
    TimestampedModel.metadata,
    Column("prompt_id", String, ForeignKey("prompts.id")),
    Column("tool_id", String, ForeignKey("tools.id")),
)


class Tool(TimestampedModel):
    """
    Defines a tool used in a chat completion. Only function calling tools are
    supported.
    """

    __tablename__ = "tools"

    function_spec: Mapped[dict] = mapped_column(JSON)


class ToolCall(TimestampedModel):
    """
    Defines a tool call included in a chat completion.
    """

    __tablename__ = "tool_calls"

    completion_id: Mapped[str] = mapped_column(ForeignKey("chat_completions.id"))
    completion: Mapped["ChatCompletion"] = relationship(
        "ChatCompletion", back_populates="tool_calls"
    )
    tool_id: Mapped[str] = mapped_column(String)
    function_name: Mapped[str] = mapped_column(String)
    function_arguments: Mapped[Optional[dict]] = mapped_column(JSON, nullable=True)

    def to_dict(self):
        """
        Convert the tool call to a dictionary format.
        """

        return remove_empty_values(
            {
                "id": self.tool_id,
                "type": "function",
                "function": {
                    "name": self.function_name,
                    "arguments": self.function_arguments,
                },
            }
        )


class Prompt(TimestampedModel):
    """
    Defines all the componenets of a prompt that was passed to a model.
    """

    __tablename__ = "prompts"

    model_id: Mapped[str] = mapped_column(ForeignKey("models.id"))
    model: Mapped["Model"] = relationship("Model", back_populates="prompts")
    messages: Mapped[List[Message]] = relationship("Message", back_populates="prompt")
    tools: Mapped[Optional[List[Tool]]] = relationship(secondary=prompt_tool_table)
    generation_kwargs: Mapped[Optional[dict]] = mapped_column(JSON, nullable=True)
    choices: Mapped[List[ChatCompletion]] = relationship(
        "ChatCompletion", back_populates="prompt"
    )
    inference_seconds: Mapped[float] = mapped_column(Float)


class TextCompletion(TimestampedModel):
    """
    Defines a text completion call with a prompt and a response.
    """

    __tablename__ = "text_completions"

    prompt: Mapped[str] = mapped_column(String)
    response: Mapped[str] = mapped_column(String)
    model: Mapped[Optional[str]] = mapped_column(String, nullable=True)

    def to_dict(self):
        """
        Convert the text completion to a dictionary format.
        """

        return remove_empty_values(
            {
                "id": self.id,
                "namespace": self.namespace,
                "prompt": self.prompt,
                "response": self.response,
                "model": self.model,
                "inference_on": self.created_at.timestamp(),
            }
        )
