"""Add name_pinyin and name_pinyin_abbr columns to identity_persons.

Revision ID: 004
Revises: 003
Create Date: 2026-03-24
"""

import sqlalchemy as sa

from alembic import op

revision = "004"
down_revision = "003"
branch_labels = None
depends_on = None


def _column_exists(table: str, column: str) -> bool:
    conn = op.get_bind()
    result = conn.execute(
        sa.text(
            "SELECT COUNT(*) FROM information_schema.columns "
            "WHERE table_schema = DATABASE() "
            "  AND table_name = :table "
            "  AND column_name = :column",
        ),
        {"table": table, "column": column},
    )
    return result.scalar() > 0


def upgrade() -> None:
    if not _column_exists("identity_persons", "name_pinyin"):
        op.add_column("identity_persons", sa.Column("name_pinyin", sa.String(256), nullable=True))
    if not _column_exists("identity_persons", "name_pinyin_abbr"):
        op.add_column("identity_persons", sa.Column("name_pinyin_abbr", sa.String(64), nullable=True))

    # Backfill pinyin for existing records
    from app.core.pinyin_utils import to_pinyin, to_pinyin_abbr

    conn = op.get_bind()
    rows = conn.execute(sa.text("SELECT id, display_name FROM identity_persons WHERE display_name IS NOT NULL"))
    for row in rows:
        pinyin = to_pinyin(row[1])
        abbr = to_pinyin_abbr(row[1])
        conn.execute(
            sa.text("UPDATE identity_persons SET name_pinyin = :pinyin, name_pinyin_abbr = :abbr WHERE id = :id"),
            {"pinyin": pinyin, "abbr": abbr, "id": row[0]},
        )


def downgrade() -> None:
    op.drop_column("identity_persons", "name_pinyin_abbr")
    op.drop_column("identity_persons", "name_pinyin")