"""Folders, chats and messages."""

from __future__ import annotations

from datetime import datetime
from typing import TYPE_CHECKING, Any

from sqlalchemy import Boolean, DateTime, ForeignKey, Integer, String, Text
from sqlalchemy.orm import Mapped, mapped_column, relationship

from lembas.db.base import Base, Timestamps, UUIDPrimaryKey
from lembas.db.types import JSONDict, JSONList

if TYPE_CHECKING:
    # Annotation only; SQLAlchemy resolves the name through its own registry at
    # runtime, so there is no import cycle. A bare `Mapped[list]` would be read
    # as a scalar and hand back None instead of [].
    from lembas.db.models.library import KnowledgeBase

ROLE_SYSTEM = "system"
ROLE_USER = "user"
ROLE_ASSISTANT = "assistant"
ROLE_TOOL = "tool"

# What a conversation is allowed to be. A plain chat can never act; an agent
# chat is pointed at a machine before it starts and stays pointed there.
KIND_CHAT = "chat"
KIND_AGENT = "agent"
KINDS = (KIND_CHAT, KIND_AGENT)

# Duplicated from services/agent/policy.py rather than imported: a model module
# importing a service would invert the dependency, and this is only the column
# default. policy.MODES is the vocabulary; this is what a row starts as.
MODE_MANUAL = "manual"


class Folder(UUIDPrimaryKey, Timestamps, Base):
    """A user-owned, arbitrarily nested container for chats."""

    __tablename__ = "folders"

    user_id: Mapped[str] = mapped_column(
        String(32), ForeignKey("users.id", ondelete="CASCADE"), nullable=False, index=True
    )
    parent_id: Mapped[str | None] = mapped_column(
        String(32), ForeignKey("folders.id", ondelete="CASCADE")
    )
    name: Mapped[str] = mapped_column(String(200), nullable=False)
    position: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
    collapsed: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)

    children: Mapped[list[Folder]] = relationship(
        back_populates="parent",
        cascade="all, delete-orphan",
        order_by="Folder.position, Folder.name",
    )
    parent: Mapped[Folder | None] = relationship(back_populates="children", remote_side="Folder.id")
    chats: Mapped[list[Chat]] = relationship(back_populates="folder")

    @property
    def visible_chats(self) -> list[Chat]:
        """The chats in this folder that belong in the sidebar.

        The relationship itself stays unfiltered -- back-population needs every
        row -- so the listing rule lives here rather than in the template, where
        the loop and the "Empty" check would have to agree by hand and already
        did not: archived chats have been showing inside folders since folders
        existed. The unfiled list has always filtered them (api/pages.py); the
        folder branch went through the relationship and filtered nothing.

        Ordered like the unfiled list: pinned first, then most recently touched.
        """
        kept = [chat for chat in self.chats if not chat.archived and not chat.temporary]
        kept.sort(key=lambda chat: chat.updated_at, reverse=True)
        kept.sort(key=lambda chat: not chat.pinned)
        return kept

    def __repr__(self) -> str:
        return f"<Folder {self.name}>"


class Chat(UUIDPrimaryKey, Timestamps, Base):
    __tablename__ = "chats"

    user_id: Mapped[str] = mapped_column(
        String(32), ForeignKey("users.id", ondelete="CASCADE"), nullable=False, index=True
    )
    # Deleting a folder keeps its chats; they fall back to the unfiled list.
    folder_id: Mapped[str | None] = mapped_column(
        String(32), ForeignKey("folders.id", ondelete="SET NULL"), index=True
    )

    title: Mapped[str] = mapped_column(String(300), default="New chat")
    # Set once the model writes the first reply, so auto-titling only runs once.
    title_generated: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)

    # Denormalised rather than a foreign key: chat history must survive an admin
    # deleting a connection or a model disappearing upstream.
    model_id: Mapped[str] = mapped_column(String(300), default="")
    connection_id: Mapped[str | None] = mapped_column(
        String(32), ForeignKey("connections.id", ondelete="SET NULL")
    )

    system_prompt: Mapped[str] = mapped_column(Text, default="")
    params_json: Mapped[dict[str, Any]] = mapped_column(JSONDict, default=dict)

    pinned: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)
    archived: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)

    # Never listed in the sidebar, and swept a day after the last thing said in
    # it. A real row rather than something held in the browser, so a reload or a
    # dropped connection does not lose the conversation -- and `Keep` clears the
    # flag, because a temporary chat that turns out to matter must have a way
    # out. See services/chat.py:sweep_temporary.
    temporary: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)

    # A reply landed while nobody was watching this chat. Cleared when the chat
    # is next opened. `unread_notified` stops the same arrival being announced
    # on every poll.
    unread: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)
    unread_notified: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)

    # --- Agent chats ---------------------------------------------------------
    # Whether this conversation may act, and where. Chosen on the new-chat
    # screen and fixed once there is a message: the harness, the tools offered
    # and the approval loop all differ, so a chat that changed kind halfway
    # would have a transcript whose earlier turns were produced under other
    # rules. The connection is locked with it -- a shell history and a project
    # directory do not transplant to another machine.
    kind: Mapped[str] = mapped_column(String(16), default=KIND_CHAT, nullable=False)
    # A plain id rather than a ForeignKey, for the reason `compacted_through_id`
    # below gives: migrations.py compiles only the column type, so a REFERENCES
    # clause would exist on a fresh database and not on an upgraded one.
    # Validated on read instead.
    ssh_profile_id: Mapped[str | None] = mapped_column(String(32))
    # Where commands start on the far side, and what file paths resolve against.
    project_dir: Mapped[str] = mapped_column(String(500), default="")
    # Which of the four permission modes is in force. The one agent field that
    # IS switchable mid-chat: it decides what gets asked about, not what the
    # conversation is.
    agent_mode: Mapped[str] = mapped_column(String(16), default=MODE_MANUAL, nullable=False)
    # Set when a turn was edited or regenerated in an agent chat. The project
    # directory is deliberately NOT rewound with the transcript -- it is
    # somebody's real working tree and deleting their work would be far worse
    # than an inconsistency -- so the harness says so instead.
    rewound_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))

    # --- Compaction ----------------------------------------------------------
    # A summary of the turns up to `compacted_through_id`, sent in their place.
    # The messages themselves are kept and still shown; they simply stop being
    # part of the request. See services/compaction.py.
    compact_summary: Mapped[str] = mapped_column(Text, default="")
    # A plain id, deliberately not a ForeignKey: db/migrations.py compiles only
    # the column type, so a REFERENCES clause would exist on a freshly created
    # database and not on an upgraded one, and a constraint half the fleet has
    # is worse than none. It is validated on every read instead -- the same
    # reasoning `model_id` above carries.
    compacted_through_id: Mapped[str | None] = mapped_column(String(32))
    compacted_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))

    folder: Mapped[Folder | None] = relationship(back_populates="chats")
    messages: Mapped[list[Message]] = relationship(
        back_populates="chat",
        cascade="all, delete-orphan",
        order_by="Message.created_at",
    )
    # Which knowledge bases this chat draws on. None means "everything its owner
    # can see"; naming some scopes the knowledge tool to those.
    knowledge_bases: Mapped[list[KnowledgeBase]] = relationship(
        "KnowledgeBase", secondary="chat_knowledge_bases"
    )

    def __repr__(self) -> str:
        return f"<Chat {self.title!r}>"


class Message(UUIDPrimaryKey, Timestamps, Base):
    __tablename__ = "messages"

    chat_id: Mapped[str] = mapped_column(
        String(32), ForeignKey("chats.id", ondelete="CASCADE"), nullable=False, index=True
    )

    # Reserved for conversation branching (edit a message, regenerate a reply
    # and keep both). Nothing reads it yet; it exists now because retrofitting a
    # column onto a live SQLite database without migrations is painful.
    parent_id: Mapped[str | None] = mapped_column(String(32), ForeignKey("messages.id"))

    role: Mapped[str] = mapped_column(String(16), nullable=False)
    content: Mapped[str] = mapped_column(Text, default="")

    # Reserved for multimodal turns: [{"type": "image_url", ...}, ...].
    # Plain-text messages leave this empty and use `content`.
    content_parts_json: Mapped[list[Any]] = mapped_column(JSONList, default=list)

    # A reasoning model's visible thinking, kept separate from the answer so it
    # can be collapsed, and so it is never fed back as context on the next turn
    # -- providers expect the answer alone, and replaying the thinking both
    # wastes the window and degrades the reply.
    reasoning: Mapped[str] = mapped_column(Text, default="")
    # Milliseconds spent producing the reasoning, for the "Thought for Xs" label.
    reasoning_ms: Mapped[int] = mapped_column(Integer, default=0, nullable=False)

    model_id: Mapped[str] = mapped_column(String(300), default="")

    # What the model did before answering: one entry per tool call, with its
    # arguments and results. Shown in the transcript so the sources behind an
    # answer stay visible, and deliberately NOT replayed as context on the next
    # turn -- see services/generation.py for why.
    tool_calls_json: Mapped[list[Any]] = mapped_column(JSONList, default=list)
    usage_json: Mapped[dict[str, Any]] = mapped_column(JSONDict, default=dict)

    # A plan produced in Plan mode: {"title": str, "steps": [str, ...]}. Marked
    # on the row rather than parsed back out of the prose, so the Execute button
    # sends exactly what was proposed and not an approximation of it.
    plan_json: Mapped[dict[str, Any]] = mapped_column(JSONDict, default=dict)

    # Non-empty when generation failed. Rendered as a styled error in the
    # thread so a failed turn is never an unexplained blank bubble.
    error: Mapped[str] = mapped_column(Text, default="")
    # False while a reply is still streaming; flipped when the stream ends.
    complete: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False)
    # True when the reader pressed Stop. Distinct from `error`: the text that
    # did arrive is kept and is perfectly usable, it is just cut short.
    stopped: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)
    # Typed while a reply was still being written, and not yet handed to a
    # model. A row rather than something held in the browser: it survives a
    # restart, it is in the transcript the moment it is typed, and it can be
    # withdrawn before it is ever sent. `build_messages` skips it; delivery --
    # `generation._drain` at the end of a reply, or `_inject` between two rounds
    # of tool calls -- is the only thing that clears it.
    queued: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)

    chat: Mapped[Chat] = relationship(back_populates="messages")
    attachments: Mapped[list[Attachment]] = relationship(  # noqa: F821
        back_populates="message",
        cascade="all, delete-orphan",
        order_by="Attachment.created_at",
    )

    @property
    def images(self) -> list:
        return [a for a in self.attachments if a.is_image]

    @property
    def documents(self) -> list:
        return [a for a in self.attachments if not a.is_image]

    def __repr__(self) -> str:
        return f"<Message {self.role} {self.content[:40]!r}>"
