"""add sms password setup and celery startup recovery

Revision ID: 476b259992de
Revises: 516b84e3bf17
Create Date: 2026-06-04 14:25:03.284184
"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = '476b259992de'
down_revision: Union[str, None] = '516b84e3bf17'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # ### chat_generation_tasks 下载容灾字段 ###
    op.add_column(
        "chat_generation_tasks",
        sa.Column("download_celery_task_id", sa.String(length=160), nullable=True),
    )
    op.add_column(
        "chat_generation_tasks",
        sa.Column("download_enqueued_at", sa.DateTime(timezone=True), nullable=True),
    )
    op.add_column(
        "chat_generation_tasks",
        sa.Column("download_started_at", sa.DateTime(timezone=True), nullable=True),
    )
    op.add_column(
        "chat_generation_tasks",
        sa.Column("download_lease_until", sa.DateTime(timezone=True), nullable=True),
    )
    op.add_column(
        "chat_generation_tasks",
        sa.Column("download_next_retry_at", sa.DateTime(timezone=True), nullable=True),
    )

    # 重点：已有数据时，新增 NOT NULL 字段必须给默认值，否则 PostgreSQL 可能直接失败。
    op.add_column(
        "chat_generation_tasks",
        sa.Column(
            "download_attempt_count",
            sa.Integer(),
            nullable=False,
            server_default=sa.text("0"),
        ),
    )

    op.add_column(
        "chat_generation_tasks",
        sa.Column("download_last_error", sa.Text(), nullable=True),
    )
    op.add_column(
        "chat_generation_tasks",
        sa.Column("download_storage_date_dir", sa.String(length=16), nullable=True),
    )

    op.create_index(
        op.f("ix_chat_generation_tasks_download_celery_task_id"),
        "chat_generation_tasks",
        ["download_celery_task_id"],
        unique=False,
    )
    op.create_index(
        op.f("ix_chat_generation_tasks_download_lease_until"),
        "chat_generation_tasks",
        ["download_lease_until"],
        unique=False,
    )
    op.create_index(
        op.f("ix_chat_generation_tasks_download_next_retry_at"),
        "chat_generation_tasks",
        ["download_next_retry_at"],
        unique=False,
    )

    # 可选：如果你不想 DB 层长期保留 server_default，可以取消下面注释。
    # 但我建议保留 DB 默认值，避免后续手写 SQL 插入时 download_attempt_count 为空。
    #
    # op.alter_column(
    #     "chat_generation_tasks",
    #     "download_attempt_count",
    #     server_default=None,
    #     existing_type=sa.Integer(),
    #     existing_nullable=False,
    # )

    # ### users 短信注册 / 设置密码字段 ###
    op.add_column(
        "users",
        sa.Column("password_set_at", sa.DateTime(timezone=True), nullable=True),
    )

    # 老用户已有密码，回填 password_set_at，避免被误判成未设置密码。
    op.execute(
        """
        UPDATE users
        SET password_set_at = CURRENT_TIMESTAMP
        WHERE hashed_password IS NOT NULL
          AND hashed_password <> ''
        """
    )

    # 短信注册用户允许暂时没有密码。
    op.alter_column(
        "users",
        "hashed_password",
        existing_type=sa.VARCHAR(length=128),
        nullable=True,
    )


def downgrade() -> None:
    # 如果已经有短信注册但未设置密码的用户，直接改回 NOT NULL 会失败。
    # 这里写一个不可登录的占位值，只是保证 downgrade 能执行。
    op.execute(
        """
        UPDATE users
        SET hashed_password = 'PASSWORD_NOT_SET_AFTER_DOWNGRADE'
        WHERE hashed_password IS NULL
           OR hashed_password = ''
        """
    )

    op.alter_column(
        "users",
        "hashed_password",
        existing_type=sa.VARCHAR(length=128),
        nullable=False,
    )

    op.drop_column("users", "password_set_at")

    op.drop_index(
        op.f("ix_chat_generation_tasks_download_next_retry_at"),
        table_name="chat_generation_tasks",
    )
    op.drop_index(
        op.f("ix_chat_generation_tasks_download_lease_until"),
        table_name="chat_generation_tasks",
    )
    op.drop_index(
        op.f("ix_chat_generation_tasks_download_celery_task_id"),
        table_name="chat_generation_tasks",
    )

    op.drop_column("chat_generation_tasks", "download_storage_date_dir")
    op.drop_column("chat_generation_tasks", "download_last_error")
    op.drop_column("chat_generation_tasks", "download_attempt_count")
    op.drop_column("chat_generation_tasks", "download_next_retry_at")
    op.drop_column("chat_generation_tasks", "download_lease_until")
    op.drop_column("chat_generation_tasks", "download_started_at")
    op.drop_column("chat_generation_tasks", "download_enqueued_at")
    op.drop_column("chat_generation_tasks", "download_celery_task_id")