"""Add chat tables (clean-foundation chat schema).

Creates:
  - chat_sessions (status as Enum: active/archived/deleted)
  - chat_messages (content NOT NULL, no NULL+FK fallback,
    no 'step' message_type — steps live in their own table)
  - chat_progress_steps (per-research step rows, separate
    sequence space from chat_messages)

Adds to research_history:
  - chat_session_id (nullable, ORM-enforced FK)
  - step_count (atomic counter for ChatService step seq)

This migration is fresh-install only. If you have a pre-release
development database with an earlier chat-table shape (legacy CHECK
constraint, nullable content, or step rows stored in chat_messages),
drop the chat tables manually before running:

    DROP TABLE IF EXISTS chat_progress_steps;
    DROP TABLE IF EXISTS chat_messages;
    DROP TABLE IF EXISTS chat_sessions;

Or simply delete the per-user encrypted DB; the auth layer
recreates it via Base.metadata.create_all on next login.

Downgrade is not supported (SQLite ALTER TABLE limits against
the legacy research_history shape).

Revision ID: 0010
Revises: 0009
Create Date: 2026-05-02
"""

from alembic import op
import sqlalchemy as sa
from loguru import logger
from sqlalchemy import inspect
from sqlalchemy_utc import UtcDateTime

# revision identifiers, used by Alembic.
revision = "0010"
down_revision = "0009"
branch_labels = None
depends_on = None


# Indexes to create for the chat tables.
#
# The composite (session_id, created_at) indexes serve the load-older
# pagination query in chat/service.py::get_session_messages, which
# filters by session_id and orders by created_at DESC. Without the
# composite, SQLite uses the single-column session_id index and sorts
# in memory; that becomes visible past ~500 rows per session. SQLite
# indexes are bidirectional, so a plain (session_id, created_at)
# satisfies both ASC and DESC ordering equally well — no DESC qualifier
# needed in the DDL.
CHAT_INDEXES = [
    ("idx_chat_session_status", "chat_sessions", ["status"]),
    ("idx_chat_session_created", "chat_sessions", ["created_at"]),
    # Composite for the sidebar list_sessions hot query, which filters by
    # status and orders by created_at DESC. With only the two single-column
    # indexes, SQLite picks one and does an in-memory sort over the result —
    # noticeable past a few hundred sessions. Same pattern as the
    # (status, created_at) composite on research_history added in 0003.
    (
        "idx_chat_session_status_created",
        "chat_sessions",
        ["status", "created_at"],
    ),
    ("ix_chat_messages_session_id", "chat_messages", ["session_id"]),
    ("ix_chat_messages_research_id", "chat_messages", ["research_id"]),
    (
        "ix_chat_messages_session_created",
        "chat_messages",
        ["session_id", "created_at"],
    ),
    (
        "ix_chat_progress_steps_research_id",
        "chat_progress_steps",
        ["research_id"],
    ),
    (
        "ix_chat_progress_steps_session_id",
        "chat_progress_steps",
        ["session_id"],
    ),
    (
        "ix_chat_progress_steps_session_created",
        "chat_progress_steps",
        ["session_id", "created_at"],
    ),
]


VALID_STATUSES = ("active", "archived", "deleted")


def _index_exists(index_name: str, table_name: str) -> bool:
    bind = op.get_bind()
    inspector = inspect(bind)
    if not inspector.has_table(table_name):
        return False
    return any(
        idx["name"] == index_name for idx in inspector.get_indexes(table_name)
    )


def _column_exists(table_name: str, column_name: str) -> bool:
    bind = op.get_bind()
    inspector = inspect(bind)
    if not inspector.has_table(table_name):
        return False
    return column_name in {
        col["name"] for col in inspector.get_columns(table_name)
    }


def _fk_exists(table_name: str, fk_name: str) -> bool:
    bind = op.get_bind()
    inspector = inspect(bind)
    if not inspector.has_table(table_name):
        return False
    return any(
        fk.get("name") == fk_name
        for fk in inspector.get_foreign_keys(table_name)
    )


def upgrade():
    bind = op.get_bind()
    inspector = inspect(bind)

    # --- chat_sessions table ---
    if not inspector.has_table("chat_sessions"):
        op.create_table(
            "chat_sessions",
            sa.Column("id", sa.String(36), primary_key=True),
            sa.Column("title", sa.String(500), nullable=True),
            sa.Column(
                "status",
                sa.Enum(*VALID_STATUSES, name="chat_session_status"),
                nullable=False,
                server_default="active",
            ),
            sa.Column("accumulated_context", sa.JSON(), nullable=True),
            sa.Column(
                "message_count",
                sa.Integer(),
                nullable=False,
                server_default="0",
            ),
            sa.Column("created_at", UtcDateTime(), nullable=False),
        )
        logger.info("0010: created chat_sessions")

    # --- chat_messages table ---
    if not inspector.has_table("chat_messages"):
        op.create_table(
            "chat_messages",
            sa.Column("id", sa.String(36), primary_key=True),
            sa.Column(
                "session_id",
                sa.String(36),
                sa.ForeignKey("chat_sessions.id", ondelete="CASCADE"),
                nullable=False,
            ),
            sa.Column(
                "research_id",
                sa.String(36),
                sa.ForeignKey("research_history.id", ondelete="SET NULL"),
                nullable=True,
            ),
            sa.Column(
                "role",
                sa.Enum("user", "assistant", name="chat_role"),
                nullable=False,
            ),
            sa.Column(
                "message_type",
                sa.Enum(
                    "query", "followup", "response", name="chat_message_type"
                ),
                nullable=False,
            ),
            sa.Column("content", sa.Text(), nullable=False),
            sa.Column("sequence_number", sa.Integer(), nullable=False),
            sa.Column("created_at", UtcDateTime(), nullable=False),
            sa.UniqueConstraint(
                "session_id",
                "sequence_number",
                name="uq_chat_message_session_seq",
            ),
        )
        logger.info("0010: created chat_messages")

    # --- chat_progress_steps table ---
    if not inspector.has_table("chat_progress_steps"):
        op.create_table(
            "chat_progress_steps",
            sa.Column("id", sa.String(36), primary_key=True),
            sa.Column(
                "research_id",
                sa.String(36),
                sa.ForeignKey("research_history.id", ondelete="CASCADE"),
                nullable=False,
            ),
            sa.Column(
                "session_id",
                sa.String(36),
                sa.ForeignKey("chat_sessions.id", ondelete="CASCADE"),
                nullable=False,
            ),
            sa.Column("phase", sa.String(64), nullable=True),
            sa.Column("content", sa.Text(), nullable=False),
            sa.Column("sequence_number", sa.Integer(), nullable=False),
            sa.Column("created_at", UtcDateTime(), nullable=False),
            sa.UniqueConstraint(
                "research_id",
                "sequence_number",
                name="uq_chat_progress_step_research_seq",
            ),
        )
        logger.info("0010: created chat_progress_steps")

    # --- Indexes (idempotent — guarded) ---
    for index_name, table_name, columns in CHAT_INDEXES:
        if not _index_exists(index_name, table_name):
            op.create_index(index_name, table_name, columns)

    # --- research_history additions ---
    # Schema source of truth is Base.metadata: 0001's create_all picks up
    # the FK on research_history.chat_session_id from the model, so on
    # fresh-install DBs the FK is already present and the guards below
    # all no-op. This block exists for any DB stamped at 0008 that
    # bypasses create_all (theoretical, but cheap to keep correct).
    #
    # Idempotency: column / index / FK checks are independent so a
    # mid-migration crash followed by re-run leaves no gaps.
    #
    # FK is added separately via batch_alter_table because Alembic's
    # SQLite dialect raises NotImplementedError on inline ForeignKey
    # in op.add_column ("No support for ALTER of constraints").
    #
    # Outer guard on `has_table("research_history")`: minimal fixtures
    # used by some regression tests stamp the alembic version at a
    # pre-0001 point with a hand-crafted DB that doesn't include
    # research_history. Real fresh-install and upgrade paths both have
    # it, so the guard only no-ops on those synthetic fixtures.
    if inspector.has_table("research_history"):
        if not _column_exists("research_history", "chat_session_id"):
            op.add_column(
                "research_history",
                sa.Column("chat_session_id", sa.String(36), nullable=True),
            )
            logger.info("0010: added research_history.chat_session_id")
        if not _index_exists(
            "ix_research_history_chat_session_id", "research_history"
        ):
            op.create_index(
                "ix_research_history_chat_session_id",
                "research_history",
                ["chat_session_id"],
            )
        if not _fk_exists(
            "research_history", "fk_research_history_chat_session_id"
        ):
            with op.batch_alter_table(
                "research_history", schema=None
            ) as batch_op:
                batch_op.create_foreign_key(
                    "fk_research_history_chat_session_id",
                    "chat_sessions",
                    ["chat_session_id"],
                    ["id"],
                    ondelete="SET NULL",
                )
            logger.info(
                "0010: added FK research_history.chat_session_id -> "
                "chat_sessions.id"
            )

        if not _column_exists("research_history", "step_count"):
            op.add_column(
                "research_history",
                sa.Column(
                    "step_count",
                    sa.Integer(),
                    nullable=False,
                    server_default="0",
                ),
            )
            logger.info("0010: added research_history.step_count")
    else:
        logger.info(
            "0010: research_history missing — skipping chat-aware column/FK/index additions"
        )

    # Partial unique index closing the SELECT-then-INSERT race in
    # chat/routes.py::send_message. The "at-most-one-in-progress per
    # chat_session_id" invariant is enforced at the DB so that two
    # near-simultaneous POSTs cannot both pass the pre-flight check
    # and both insert IN_PROGRESS rows. With the index in place, the
    # second concurrent INSERT raises IntegrityError and the existing
    # `except IntegrityError` handler in routes.py converts that into
    # HTTP 409 — making the race truly atomic.
    #
    # Index is intentionally partial:
    # * status='in_progress' — only the in-flight slot is unique;
    #   arbitrarily many completed/failed/terminated runs per chat
    #   session must remain allowed (history view).
    # * chat_session_id IS NOT NULL — non-chat research (news,
    #   scheduler, direct API) doesn't carry chat_session_id and
    #   must not be constrained.
    #
    # SQLite supports partial indexes since 3.8.0 (2014); project
    # floor is Python >= 3.12 which bundles SQLite >= 3.39.
    _idx_name = "ux_research_history_chat_session_in_progress"
    if inspector.has_table("research_history") and not _index_exists(
        _idx_name, "research_history"
    ):
        op.create_index(
            _idx_name,
            "research_history",
            ["chat_session_id"],
            unique=True,
            sqlite_where=sa.text(
                "status = 'in_progress' AND chat_session_id IS NOT NULL"
            ),
            postgresql_where=sa.text(
                "status = 'in_progress' AND chat_session_id IS NOT NULL"
            ),
        )
        logger.info(
            "0010: added partial unique index "
            "ux_research_history_chat_session_in_progress"
        )


def downgrade():
    """Downgrade not supported (NotImplementedError).

    Why: SQLite ALTER TABLE has a hard limitation against dropping a
    column that's the target of a FOREIGN KEY definition, even with
    PRAGMA foreign_keys=OFF. The model layer's create_all path adds
    a FK on research_history.chat_session_id that this migration's
    op.add_column does NOT add — leaving fresh-install DBs and
    upgrade-built DBs in different shapes for downgrade purposes.
    Alembic's batch_alter_table fails on the legacy research_history
    schema's unnamed constraints ("Constraint must have a name"),
    and a hand-rolled CREATE-TABLE-and-copy approach is fragile to
    parse around the FK definitions.

    The project is dev-stage with no live users; the supported
    rollback is to recreate the per-user database. The pre-migration
    backup created by the encrypted-DB backup step is the recovery
    artifact.

    Tests in test_alembic_migrations.py that exercise downgrade
    parametrize against the migration chain — those tests skip
    revision "0010" via the `NON_REVERSIBLE_REVISIONS` exemption
    (mirroring the existing `INTENTIONAL_NOOP_DOWNGRADES` pattern).
    """
    raise NotImplementedError(
        "0010 (chat tables) is not reversible due to SQLite ALTER TABLE "
        "limitations against the legacy research_history shape. "
        "Recreate the dev database to roll back, or restore from the "
        "pre-migration backup."
    )
