"""Add tier, daily_translation_count, and hashed_password fields

Revision ID: 002_add_tier_daily_count
Revises: 001_initial
Create Date: 2026-02-20

"""

from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


revision: str = "002_add_tier_daily_count"
down_revision: Union[str, None] = "001_initial"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    conn = op.get_bind()
    dialect = conn.dialect.name

    op.add_column(
        "users",
        sa.Column("tier", sa.String(10), nullable=True, server_default="free"),
    )
    op.add_column(
        "users",
        sa.Column(
            "daily_translation_count",
            sa.Integer(),
            nullable=True,
            server_default=sa.text("0"),
        ),
    )
    op.add_column("users", sa.Column("hashed_password", sa.String(255), nullable=True))

    op.execute(
        sa.text(
            "UPDATE users SET hashed_password = password_hash WHERE hashed_password IS NULL"
        )
    )

    op.execute(
        sa.text(
            "UPDATE users SET tier = CASE "
            "WHEN plan IN ('pro', 'business', 'enterprise') THEN 'pro' "
            "ELSE 'free' END "
            "WHERE tier IS NULL OR tier = 'free'"
        )
    )

    op.execute(
        sa.text(
            "UPDATE users SET daily_translation_count = 0 WHERE daily_translation_count IS NULL"
        )
    )

    # SQLite does not support ALTER COLUMN — skip nullable enforcement there
    if dialect != "sqlite":
        op.alter_column("users", "tier", nullable=False)
        op.alter_column("users", "daily_translation_count", nullable=False)
        op.create_check_constraint("ck_users_tier", "users", "tier IN ('free', 'pro')")


def downgrade() -> None:
    conn = op.get_bind()
    dialect = conn.dialect.name

    if dialect != "sqlite":
        op.drop_constraint("ck_users_tier", "users", type_="check")

    op.drop_column("users", "hashed_password")
    op.drop_column("users", "daily_translation_count")
    op.drop_column("users", "tier")