"""add snapshot fields and check constraints

Revision ID: 0004_snapshot_and_constraints
Revises: 0003_injuries
Create Date: 2026-09-15

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa

# revision identifiers, used by Alembic.
revision: str = '0004_snapshot_and_constraints'
down_revision: Union[str, None] = '0003_injuries'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # 1. 重命名 confidence → subjective_confidence
    op.alter_column('predictions', 'confidence', new_column_name='subjective_confidence')

    # 2. 新增快照字段
    op.add_column('predictions', sa.Column('cutoff_at', sa.DateTime(timezone=True), nullable=True))
    op.add_column('predictions', sa.Column('input_hash', sa.String(length=64), nullable=True))

    # 3. 新增 CHECK 约束
    #    注意: 对已有数据行加 CHECK 约束时,若存在越界数据 ALTER TABLE 会中途失败,
    #    导致迁移卡在半完成状态。这里先做一次性清洗(把越界值收敛到合法域),
    #    再建约束,保证在非空库上也能成功。
    op.execute("UPDATE predictions SET pred_home_goals = 0 WHERE pred_home_goals IS NOT NULL AND pred_home_goals < 0")
    op.execute("UPDATE predictions SET pred_away_goals = 0 WHERE pred_away_goals IS NOT NULL AND pred_away_goals < 0")
    op.execute(
        "UPDATE predictions SET subjective_confidence = "
        "CASE WHEN subjective_confidence < 0 THEN 0 "
        "     WHEN subjective_confidence > 1 THEN 1 ELSE subjective_confidence END "
        "WHERE subjective_confidence IS NOT NULL "
        "  AND (subjective_confidence < 0 OR subjective_confidence > 1)"
    )
    op.execute("UPDATE predictions SET pred_1x2 = NULL WHERE pred_1x2 IS NOT NULL AND pred_1x2 NOT IN ('1', 'X', '2')")
    op.execute("UPDATE predictions SET mode = 'single' WHERE mode IS NOT NULL AND mode NOT IN ('single', 'multi')")

    op.create_check_constraint('ck_pred_home_goals_nonneg', 'predictions', 'pred_home_goals >= 0')
    op.create_check_constraint('ck_pred_away_goals_nonneg', 'predictions', 'pred_away_goals >= 0')
    op.create_check_constraint('ck_confidence_range', 'predictions', 'subjective_confidence >= 0 AND subjective_confidence <= 1')
    op.create_check_constraint('ck_pred_1x2_enum', 'predictions', "pred_1x2 IN ('1', 'X', '2')")
    op.create_check_constraint('ck_mode_enum', 'predictions', "mode IN ('single', 'multi')")


def downgrade() -> None:
    # 1. 删除 CHECK 约束
    op.drop_constraint('ck_mode_enum', 'predictions', type_='check')
    op.drop_constraint('ck_pred_1x2_enum', 'predictions', type_='check')
    op.drop_constraint('ck_confidence_range', 'predictions', type_='check')
    op.drop_constraint('ck_pred_away_goals_nonneg', 'predictions', type_='check')
    op.drop_constraint('ck_pred_home_goals_nonneg', 'predictions', type_='check')

    # 2. 删除快照字段
    op.drop_column('predictions', 'input_hash')
    op.drop_column('predictions', 'cutoff_at')

    # 3. 恢复列名
    op.alter_column('predictions', 'subjective_confidence', new_column_name='confidence')