"""segments：Phase2 低置信路由延迟元数据

Revision ID: 0017_segment_narrative_defer
Revises: 0016_memory_pipeline_status
"""

from typing import Sequence, Union

import sqlalchemy as sa

from alembic import op

revision: str = "0017_segment_narrative_defer"
down_revision: Union[str, None] = "0016_memory_pipeline_status"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def _column_names(table_name: str) -> set[str]:
    bind = op.get_bind()
    inspector = sa.inspect(bind)
    return {column["name"] for column in inspector.get_columns(table_name)}


def upgrade() -> None:
    columns = _column_names("segments")
    if "narrative_deferred_until" not in columns:
        op.add_column(
            "segments",
            sa.Column(
                "narrative_deferred_until",
                sa.DateTime(timezone=True),
                nullable=True,
            ),
        )
    if "narrative_defer_count" not in columns:
        op.add_column(
            "segments",
            sa.Column(
                "narrative_defer_count",
                sa.Integer(),
                nullable=False,
                server_default=sa.text("0"),
            ),
        )
    if "narrative_defer_reason" not in columns:
        op.add_column(
            "segments",
            sa.Column(
                "narrative_defer_reason",
                sa.String(),
                nullable=True,
            ),
        )
    if "narrative_last_attempt_at" not in columns:
        op.add_column(
            "segments",
            sa.Column(
                "narrative_last_attempt_at",
                sa.DateTime(timezone=True),
                nullable=True,
            ),
        )


def downgrade() -> None:
    columns = _column_names("segments")
    for column in (
        "narrative_last_attempt_at",
        "narrative_defer_reason",
        "narrative_defer_count",
        "narrative_deferred_until",
    ):
        if column in columns:
            op.drop_column("segments", column)