"""Memory pipeline status columns.

Revision ID: 0016_memory_pipeline_status
Revises: 0015_memory_single_chain
"""

from typing import Sequence, Union

import sqlalchemy as sa

from alembic import op

revision: str = "0016_memory_pipeline_status"
down_revision: Union[str, None] = "0015_memory_single_chain"
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("memory_sources", "embedding_status"):
        op.add_column(
            "memory_sources",
            sa.Column("embedding_status", sa.String(), nullable=True),
        )
    if not _has_column("memory_sources", "embedding_error"):
        op.add_column(
            "memory_sources",
            sa.Column("embedding_error", sa.Text(), nullable=True),
        )
    if not _has_column("memory_sources", "enrichment_status"):
        op.add_column(
            "memory_sources",
            sa.Column("enrichment_status", sa.String(), nullable=True),
        )
    if not _has_column("memory_sources", "enrichment_error"):
        op.add_column(
            "memory_sources",
            sa.Column("enrichment_error", sa.Text(), nullable=True),
        )
    if not _has_column("memory_chunks", "embedding_status"):
        op.add_column(
            "memory_chunks",
            sa.Column("embedding_status", sa.String(), nullable=True),
        )
    if not _has_column("memory_chunks", "embedding_error"):
        op.add_column(
            "memory_chunks",
            sa.Column("embedding_error", sa.Text(), nullable=True),
        )

    op.execute(
        """
        UPDATE memory_sources
        SET
          embedding_status = COALESCE(embedding_status, 'success'),
          enrichment_status = COALESCE(enrichment_status, 'pending')
        """
    )
    op.execute(
        """
        UPDATE memory_chunks
        SET embedding_status = COALESCE(
          embedding_status,
          CASE WHEN embedding IS NULL THEN 'pending' ELSE 'success' END
        )
        """
    )
    op.create_index(
        "ix_memory_sources_user_embedding_status",
        "memory_sources",
        ["user_id", "embedding_status"],
        unique=False,
    )
    op.create_index(
        "ix_memory_sources_user_enrichment_status",
        "memory_sources",
        ["user_id", "enrichment_status"],
        unique=False,
    )
    op.create_index(
        "ix_memory_chunks_user_embedding_status",
        "memory_chunks",
        ["user_id", "embedding_status"],
        unique=False,
    )


def downgrade() -> None:
    op.drop_index("ix_memory_chunks_user_embedding_status", table_name="memory_chunks")
    op.drop_index(
        "ix_memory_sources_user_enrichment_status", table_name="memory_sources"
    )
    op.drop_index("ix_memory_sources_user_embedding_status", table_name="memory_sources")
    for column in ("embedding_error", "embedding_status"):
        if _has_column("memory_chunks", column):
            op.drop_column("memory_chunks", column)
    for column in (
        "enrichment_error",
        "enrichment_status",
        "embedding_error",
        "embedding_status",
    ):
        if _has_column("memory_sources", column):
            op.drop_column("memory_sources", column)