"""Add token_hash columns to oauth_grants and oauth_codes

Revision ID: acct_0002
Revises: acct_0001
Create Date: 2026-02-26
"""

import hashlib
import sqlalchemy as sa
from alembic import op

revision = "acct_0002"
down_revision = "acct_0001"
branch_labels = None
depends_on = None


def _hash(token: str) -> str:
    return hashlib.sha256(token.encode()).hexdigest()


def upgrade():
    # Add new hash columns
    op.add_column("oauth_grants", sa.Column("token_hash", sa.String(64), nullable=True))
    op.add_column("oauth_codes", sa.Column("code_hash", sa.String(64), nullable=True))
    op.add_column("oauth_codes", sa.Column("grant_token_hash", sa.String(64), nullable=True))

    # Backfill hashes from existing plaintext tokens
    conn = op.get_bind()
    grants = conn.execute(sa.text("SELECT id, token FROM oauth_grants WHERE token IS NOT NULL"))
    for row in grants:
        conn.execute(
            sa.text("UPDATE oauth_grants SET token_hash = :h WHERE id = :id"),
            {"h": _hash(row.token), "id": row.id},
        )

    codes = conn.execute(sa.text("SELECT id, code, grant_token FROM oauth_codes WHERE code IS NOT NULL"))
    for row in codes:
        params = {"id": row.id, "ch": _hash(row.code)}
        params["gh"] = _hash(row.grant_token) if row.grant_token else None
        conn.execute(
            sa.text("UPDATE oauth_codes SET code_hash = :ch, grant_token_hash = :gh WHERE id = :id"),
            params,
        )

    # Create unique indexes on hash columns
    op.create_index("ix_oauth_grant_token_hash", "oauth_grants", ["token_hash"], unique=True)
    op.create_index("ix_oauth_code_code_hash", "oauth_codes", ["code_hash"], unique=True)

    # Make original token columns nullable (keep for rollback safety)
    op.alter_column("oauth_grants", "token", nullable=True)
    op.alter_column("oauth_codes", "code", nullable=True)

    # Drop old unique indexes on plaintext columns
    try:
        op.drop_index("ix_oauth_grant_token", "oauth_grants")
    except Exception:
        pass
    try:
        op.drop_index("ix_oauth_code_code", "oauth_codes")
    except Exception:
        pass


def downgrade():
    # Restore original NOT NULL constraints
    op.alter_column("oauth_grants", "token", nullable=False)
    op.alter_column("oauth_codes", "code", nullable=False)

    # Drop hash columns and indexes
    try:
        op.drop_index("ix_oauth_grant_token_hash", "oauth_grants")
    except Exception:
        pass
    try:
        op.drop_index("ix_oauth_code_code_hash", "oauth_codes")
    except Exception:
        pass

    op.drop_column("oauth_grants", "token_hash")
    op.drop_column("oauth_codes", "code_hash")
    op.drop_column("oauth_codes", "grant_token_hash")

    # Restore original unique indexes
    op.create_index("ix_oauth_grant_token", "oauth_grants", ["token"], unique=True)
    op.create_index("ix_oauth_code_code", "oauth_codes", ["code"], unique=True)