"""Dialogue lineage: segment/memory/story/chapter evidence message-grade provenance.

Revision ID: 0011_dialogue_lineage
Revises: 0010_ce_snapshots
"""

from typing import Sequence, Union

import sqlalchemy as sa
from sqlalchemy.dialects import postgresql

from alembic import op

revision: str = "0011_dialogue_lineage"
down_revision: Union[str, None] = "0010_ce_snapshots"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def _has_column(table: str, column: str) -> bool:
    bind = op.get_bind()
    return any(c["name"] == column for c in sa.inspect(bind).get_columns(table))


def upgrade() -> None:
    if not _has_column("segments", "user_message_id"):
        op.add_column(
            "segments",
            sa.Column(
                "user_message_id",
                sa.String(),
                sa.ForeignKey("conversation_messages.id", ondelete="SET NULL"),
                nullable=True,
            ),
        )
        op.create_index(
            "ix_segments_user_message_id",
            "segments",
            ["user_message_id"],
            unique=False,
        )
    if not _has_column("segments", "lineage_json"):
        op.add_column(
            "segments",
            sa.Column(
                "lineage_json",
                postgresql.JSON(astext_type=sa.Text()),
                nullable=True,
            ),
        )

    if not _has_column("memory_sources", "lineage_json"):
        op.add_column(
            "memory_sources",
            sa.Column(
                "lineage_json",
                postgresql.JSON(astext_type=sa.Text()),
                nullable=True,
            ),
        )
    if not _has_column("memory_sources", "primary_user_message_id"):
        op.add_column(
            "memory_sources",
            sa.Column("primary_user_message_id", sa.String(), nullable=True),
        )
        op.create_index(
            "ix_memory_sources_primary_user_message_id",
            "memory_sources",
            ["primary_user_message_id"],
            unique=False,
        )

    if not _has_column("story_versions", "lineage_json"):
        op.add_column(
            "story_versions",
            sa.Column(
                "lineage_json",
                postgresql.JSON(astext_type=sa.Text()),
                nullable=True,
            ),
        )

    if not _has_column("chapter_evidence_snapshots", "message_lineage_json"):
        op.add_column(
            "chapter_evidence_snapshots",
            sa.Column(
                "message_lineage_json",
                postgresql.JSON(astext_type=sa.Text()),
                nullable=True,
            ),
        )

    if not _has_column("chapters", "source_lineage_json"):
        op.add_column(
            "chapters",
            sa.Column(
                "source_lineage_json",
                postgresql.JSON(astext_type=sa.Text()),
                nullable=True,
            ),
        )


def downgrade() -> None:
    if _has_column("chapters", "source_lineage_json"):
        op.drop_column("chapters", "source_lineage_json")
    if _has_column("chapter_evidence_snapshots", "message_lineage_json"):
        op.drop_column("chapter_evidence_snapshots", "message_lineage_json")
    if _has_column("story_versions", "lineage_json"):
        op.drop_column("story_versions", "lineage_json")
    if _has_column("memory_sources", "primary_user_message_id"):
        op.drop_index(
            "ix_memory_sources_primary_user_message_id", table_name="memory_sources"
        )
        op.drop_column("memory_sources", "primary_user_message_id")
    if _has_column("memory_sources", "lineage_json"):
        op.drop_column("memory_sources", "lineage_json")
    if _has_column("segments", "lineage_json"):
        op.drop_column("segments", "lineage_json")
    if _has_column("segments", "user_message_id"):
        op.drop_index("ix_segments_user_message_id", table_name="segments")
        op.drop_column("segments", "user_message_id")