from datetime import datetime, timezone
from typing import TYPE_CHECKING

from sqlalchemy import JSON, DateTime, Integer, String
from sqlalchemy.orm import Mapped, mapped_column, relationship
from uuid6 import uuid7

from vdash.core import BaseModel

if TYPE_CHECKING:
    from vdash.modules.sql import SqlQuery


class ChartConversation(BaseModel):
    __tablename__ = "chart_conversations"

    uuid: Mapped[str] = mapped_column(
        String(36),
        unique=True,
        nullable=False,
        default=lambda: str(uuid7()),
        index=True,
    )

    sql_query_id: Mapped[int] = mapped_column(
        Integer,
        nullable=False,
        index=True,
    )

    title: Mapped[str | None] = mapped_column(
        String(500),
        nullable=True,
        doc="Auto-generated title from user input",
    )

    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True),
        default=lambda: datetime.now(timezone.utc),
        nullable=False,
    )

    updated_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True),
        default=lambda: datetime.now(timezone.utc),
        onupdate=lambda: datetime.now(timezone.utc),
        nullable=False,
    )

    # Relationship (logical link, no FK constraint)
    sql_query: Mapped["SqlQuery"] = relationship(
        primaryjoin="foreign(ChartConversation.sql_query_id) == SqlQuery.id",
        lazy="joined",
    )

    messages: Mapped[list["ChartMessage"]] = relationship(
        primaryjoin="ChartConversation.id == foreign(ChartMessage.conversation_id)",
        order_by="ChartMessage.created_at",
        back_populates="conversation",
    )


class ChartMessage(BaseModel):
    __tablename__ = "chart_messages"

    uuid: Mapped[str] = mapped_column(
        String(36),
        unique=True,
        nullable=False,
        default=lambda: str(uuid7()),
        index=True,
    )

    conversation_id: Mapped[int] = mapped_column(
        Integer,
        nullable=False,
        index=True,
    )

    role: Mapped[str] = mapped_column(
        String(20),
        nullable=False,
        doc="user | assistant",
    )

    content: Mapped[dict] = mapped_column(
        JSON,
        nullable=False,
        doc="TextContent | ChartContent",
    )

    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True),
        default=lambda: datetime.now(timezone.utc),
        nullable=False,
    )

    # Relationship (logical link, no FK constraint)
    conversation: Mapped["ChartConversation"] = relationship(
        primaryjoin="foreign(ChartMessage.conversation_id) == ChartConversation.id",
        back_populates="messages",
    )

    chart: Mapped["Chart | None"] = relationship(
        primaryjoin="ChartMessage.id == foreign(Chart.message_id)",
        back_populates="message",
        uselist=False,
    )


class Chart(BaseModel):
    __tablename__ = "charts"

    uuid: Mapped[str] = mapped_column(
        String(36),
        unique=True,
        nullable=False,
        default=lambda: str(uuid7()),
        index=True,
    )

    message_id: Mapped[int] = mapped_column(
        Integer,
        nullable=False,
        index=True,
    )

    sql_query_id: Mapped[int | None] = mapped_column(
        Integer,
        nullable=True,
        index=True,
    )

    name: Mapped[str | None] = mapped_column(
        String(500),
        nullable=True,
        doc="Auto-generated name from user input",
    )

    vega_chart: Mapped[dict] = mapped_column(
        JSON,
        nullable=False,
        doc="Vega chart specification",
    )

    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True),
        default=lambda: datetime.now(timezone.utc),
        nullable=False,
    )

    updated_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True),
        default=lambda: datetime.now(timezone.utc),
        onupdate=lambda: datetime.now(timezone.utc),
        nullable=False,
    )

    # Relationship (logical link, no FK constraint)
    sql_query: Mapped["SqlQuery"] = relationship(
        primaryjoin="foreign(Chart.sql_query_id) == SqlQuery.id",
        lazy="joined",
    )

    message: Mapped["ChartMessage"] = relationship(
        primaryjoin="foreign(Chart.message_id) == ChartMessage.id",
        back_populates="chart",
    )
