"""add role column to users

Revision ID: 25c6bf382eed
Revises: 0f0c0a931d62
Create Date: 2025-05-24 13:51:36.181191

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql

# revision identifiers, used by Alembic.
revision: str = '25c6bf382eed'
down_revision: Union[str, None] = '0f0c0a931d62'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.alter_column('application_users', 'created_at',
               existing_type=postgresql.TIMESTAMP(),
               type_=sa.DateTime(timezone=True),
               nullable=False)
    op.alter_column('application_users', 'updated_at',
               existing_type=postgresql.TIMESTAMP(),
               type_=sa.DateTime(timezone=True),
               existing_nullable=True)
    plan_enum = sa.Enum('FREE', 'STARTER', 'PRO', 'ENTERPRISE', name='CompanyPlanEnum')
    plan_enum.create(op.get_bind(), checkfirst=True)

    # Explicit type cast
    op.execute("""
        ALTER TABLE companies
        ALTER COLUMN plan
        TYPE CompanyPlanEnum
        USING plan::CompanyPlanEnum
    """)
    # Create the ENUM first
    role_enum = sa.Enum('ADMIN', 'ANALYST', 'VIEWER', name='companyuserroleenum')
    role_enum.create(op.get_bind(), checkfirst=True)

    # Step 1: Add the column as nullable
    op.add_column('users', sa.Column('role', sa.Enum('ADMIN', 'ANALYST', 'VIEWER', name='companyuserroleenum'), nullable=True))

    # Step 2: Set a default role for existing rows
    op.execute("""
        UPDATE users SET role = 'VIEWER' WHERE role IS NULL
    """)

    # Step 3: Alter the column to be NOT NULL
    op.alter_column('users', 'role', nullable=False)

    department_enum = sa.Enum('HR', 'FINANCE', 'MARKETING', 'SALES', 'CEO', 'CFO', 'ADMINISTRATION', 'IT', name='departmentenum')
    department_enum.create(op.get_bind(), checkfirst=True)

    # Step 2: Cast the existing TEXT column to ENUM using raw SQL
    op.execute("""
        ALTER TABLE users
        ALTER COLUMN department
        TYPE departmentenum
        USING department::departmentenum
    """)
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.alter_column('users', 'department',
               existing_type=sa.Enum('HR', 'FINANCE', 'MARKETING', 'SALES', 'CEO', 'CFO', 'ADMINISTRATION', 'IT', name='departmentenum'),
               type_=sa.TEXT(),
               existing_nullable=False)
    op.drop_column('users', 'role')
    op.alter_column('companies', 'plan',
               existing_type=sa.Enum('FREE', 'STARTER', 'PRO', 'ENTERPRISE', name='CompanyPlanEnum'),
               type_=sa.TEXT(),
               existing_nullable=False)
    op.alter_column('application_users', 'updated_at',
               existing_type=sa.DateTime(timezone=True),
               type_=postgresql.TIMESTAMP(),
               existing_nullable=True)
    op.alter_column('application_users', 'created_at',
               existing_type=sa.DateTime(timezone=True),
               type_=postgresql.TIMESTAMP(),
               nullable=True)
    # ### end Alembic commands ###
