"""Add multi_tenancy

Revision ID: ff9052a97835
Revises: ff8052a97834
Create Date: 2026-03-05 10:00:00.000000

"""

from typing import Sequence, Union

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


# revision identifiers, used by Alembic.
revision: str = "ff9052a97835"
down_revision: Union[str, None] = "ff8052a97834"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # 1. Create tenants table
    op.execute(
        """
    CREATE TABLE tenants (
        id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
        name VARCHAR(255) NOT NULL,
        slug VARCHAR(100) NOT NULL UNIQUE,
        is_active BOOLEAN NOT NULL DEFAULT TRUE,
        created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
        updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()
    );
    """
    )


    # 3. Add tenant_id columns (nullable first)
    tables = [
        "users",
        "projects",
        "drive_folders",
        "drive_files",
        "drive_activities",
        "chatbot_documents",
        "notifications",
        "user_files",
    ]
    for table in tables:
        op.execute(
            f"ALTER TABLE {table} ADD COLUMN tenant_id UUID REFERENCES tenants(id);"
        )

    # 4. Backfill existing data
    for table in tables:
        op.execute(
            f"UPDATE {table} SET tenant_id = '00000000-0000-0000-0000-000000000001';"
        )

    # 5. Set NOT NULL (except users)
    not_null_tables = [
        "projects",
        "drive_folders",
        "drive_files",
        "drive_activities",
        "chatbot_documents",
        "notifications",
        "user_files",
    ]
    for table in not_null_tables:
        op.execute(f"ALTER TABLE {table} ALTER COLUMN tenant_id SET NOT NULL;")

    # 6. Create indexes
    for table in tables:
        op.execute(f"CREATE INDEX ix_{table}_tenant_id ON {table}(tenant_id);")

    # 7. Constraint unique (tenant_id, email)
    op.execute("ALTER TABLE users DROP CONSTRAINT IF EXISTS users_email_key;")
    op.execute("ALTER TABLE users DROP CONSTRAINT IF EXISTS uq_users_email;")
    op.execute(
        "ALTER TABLE users ADD CONSTRAINT uq_user_tenant_email UNIQUE (tenant_id, email);"
    )


def downgrade() -> None:
    op.execute("ALTER TABLE users DROP CONSTRAINT IF EXISTS uq_user_tenant_email;")
    op.execute("ALTER TABLE users ADD CONSTRAINT users_email_key UNIQUE (email);")

    tables = [
        "users",
        "projects",
        "drive_folders",
        "drive_files",
        "drive_activities",
        "chatbot_documents",
        "notifications",
        "user_files",
    ]
    for table in tables:
        op.execute(f"DROP INDEX IF EXISTS ix_{table}_tenant_id;")
        op.execute(f"ALTER TABLE {table} DROP COLUMN IF EXISTS tenant_id;")

    op.execute("DROP TABLE IF EXISTS tenants;")