"""shift all naive timestamps to beijing and set db timezone

Revision ID: shift_all_timestamps_to_beijing
Revises: add_conversation_sessions_table
Create Date: 2026-03-11
"""

from __future__ import annotations

from alembic import op
import sqlalchemy as sa


revision = "shift_all_timestamps_to_beijing"
down_revision = "add_conversation_sessions_table"
branch_labels = None
depends_on = None


def _is_postgresql() -> bool:
    bind = op.get_bind()
    return str(getattr(getattr(bind, "dialect", None), "name", "")).lower().startswith("postgres")


def _list_timestamp_columns() -> list[tuple[str, str]]:
    bind = op.get_bind()
    rows = bind.execute(
        sa.text(
            """
            SELECT table_name, column_name
            FROM information_schema.columns
            WHERE table_schema = 'public'
              AND data_type = 'timestamp without time zone'
            ORDER BY table_name, column_name
            """
        )
    ).fetchall()
    return [(str(row[0]), str(row[1])) for row in rows]


def _shift_all_timestamp_columns(*, hours: int) -> None:
    delta_expr = f"INTERVAL '{int(hours)} hours'" if int(hours) >= 0 else f"- INTERVAL '{abs(int(hours))} hours'"
    for table_name, column_name in _list_timestamp_columns():
        if int(hours) >= 0:
            sql = f'UPDATE "{table_name}" SET "{column_name}" = "{column_name}" + {delta_expr} WHERE "{column_name}" IS NOT NULL'
        else:
            sql = f'UPDATE "{table_name}" SET "{column_name}" = "{column_name}" {delta_expr} WHERE "{column_name}" IS NOT NULL'
        op.execute(sa.text(sql))


def _set_database_timezone(timezone_name: str) -> None:
    bind = op.get_bind()
    database_name = str(bind.execute(sa.text("SELECT current_database()")).scalar_one())
    role_name = str(bind.execute(sa.text("SELECT current_user")).scalar_one())
    op.execute(sa.text(f'ALTER DATABASE "{database_name}" SET timezone TO \'{timezone_name}\''))
    op.execute(sa.text(f'ALTER ROLE "{role_name}" SET timezone TO \'{timezone_name}\''))


def upgrade() -> None:
    if not _is_postgresql():
        return
    _set_database_timezone("Asia/Shanghai")
    _shift_all_timestamp_columns(hours=8)


def downgrade() -> None:
    if not _is_postgresql():
        return
    _shift_all_timestamp_columns(hours=-8)
    # 统一策略：即使降级历史数据，也保持数据库会话时区为北京时间。
    _set_database_timezone("Asia/Shanghai")