from __future__ import annotations

from datetime import UTC, datetime

from sqlalchemy import ForeignKey, select
from sqlalchemy.dialects.postgresql import JSONB, TIMESTAMP
from sqlalchemy.ext.asyncio import AsyncAttrs
from sqlalchemy.orm import (
    DeclarativeBase,
    Mapped,
    MappedAsDataclass,
    column_property,
    mapped_column,
    relationship,
)


class Base(DeclarativeBase, MappedAsDataclass, AsyncAttrs):
    type_annotation_map = {
        datetime: TIMESTAMP(timezone=True),
        dict: JSONB,
    }


class CommonMixin(MappedAsDataclass):
    id: Mapped[int] = mapped_column(init=False, primary_key=True)
    created_at: Mapped[datetime] = mapped_column(
        init=False, default_factory=lambda: datetime.now(UTC)
    )
    updated_at: Mapped[datetime | None] = mapped_column(
        init=False, default=None, onupdate=lambda: datetime.now(UTC)
    )
    active: Mapped[bool] = mapped_column(init=False, default=True)


class User(CommonMixin, Base):
    __tablename__ = "users"

    name: Mapped[str] = mapped_column(unique=True, index=True)

    conversations: Mapped[list[Conversation]] = relationship(
        init=False, back_populates="user"
    )


class Assistant(CommonMixin, Base):
    __tablename__ = "assistants"

    name: Mapped[str] = mapped_column(unique=True, index=True)
    system_prompt: Mapped[str]

    conversations: Mapped[list[Conversation]] = relationship(
        init=False, back_populates="assistant"
    )


class Conversation(CommonMixin, Base):
    __tablename__ = "conversations"

    user_id: Mapped[int] = mapped_column(ForeignKey("users.id"))
    assistant_id: Mapped[int] = mapped_column(ForeignKey("assistants.id"))

    user: Mapped[User] = relationship(init=False, back_populates="conversations")
    assistant: Mapped[Assistant] = relationship(
        init=False, back_populates="conversations"
    )
    messages: Mapped[list[Message]] = relationship(
        init=False, order_by="Message.created_at"
    )

    assistant_name = column_property(
        select(Assistant.name).where(Assistant.id == assistant_id).scalar_subquery()
    )

    def add_message(
        self,
        role: str,
        content: str | None = None,
        tools: dict | None = None,
        query_id: int | None = None,
    ):
        self.messages.append(
            Message(
                role=role,
                content=content,
                tools=tools,
                query_id=query_id,
                conversation_id=self.id,
            )
        )

    async def to_openai_format(self):
        messages = await self.awaitable_attrs.messages
        return [(await m.to_openai_format()) for m in messages]


class Message(CommonMixin, Base):
    __tablename__ = "messages"

    conversation_id: Mapped[int] = mapped_column(ForeignKey("conversations.id"))

    role: Mapped[str]
    content: Mapped[str | None] = mapped_column(default=None)
    feedback: Mapped[bool | None] = mapped_column(default=None)
    tools: Mapped[dict | None] = mapped_column(default=None)
    query_id: Mapped[int | None] = mapped_column(default=None)

    async def to_openai_format(self):
        role = await self.awaitable_attrs.role
        content = await self.awaitable_attrs.content
        tools = await self.awaitable_attrs.tools
        return {
            "role": role,
            "content": content,
            **(tools or {}),
        }


class Comment(CommonMixin, Base):
    __tablename__ = "comments"

    message_id: Mapped[int] = mapped_column(ForeignKey("messages.id"))
    content: Mapped[str]
