"""cloud_shards + project_cloud_assignments — D1 shard routing model.

Every project is assigned to a shard which owns the D1 database, R2 bucket,
and dispatch namespace for that project's data.  Shared shards serve Basic/Pro;
dedicated and customer_owned shards serve Enterprise.

Revision ID: 049
Revises: 048
"""
from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects.postgresql import ENUM, UUID as PGUUID

revision = "049"
down_revision = "048"
branch_labels = None
depends_on = None


def upgrade() -> None:
    # Idempotent: previous failed runs (or sa.Enum emitting CREATE TYPE) can leave these types
    # behind. Match pattern from migration 033.
    op.execute(
        """
        DO $$ BEGIN
            CREATE TYPE p2hs_cloud_shard_type AS ENUM ('shared', 'dedicated', 'customer_owned');
        EXCEPTION WHEN duplicate_object THEN NULL;
        END $$;
        """
    )
    op.execute(
        """
        DO $$ BEGIN
            CREATE TYPE p2hs_cloud_shard_status AS ENUM ('active', 'full', 'draining', 'offline');
        EXCEPTION WHEN duplicate_object THEN NULL;
        END $$;
        """
    )
    op.execute(
        """
        DO $$ BEGIN
            CREATE TYPE p2hs_project_shard_status AS ENUM ('active', 'migrating', 'draining');
        EXCEPTION WHEN duplicate_object THEN NULL;
        END $$;
        """
    )

    shard_type_col = ENUM(
        "shared",
        "dedicated",
        "customer_owned",
        name="p2hs_cloud_shard_type",
        create_type=False,
    )
    shard_status_col = ENUM(
        "active",
        "full",
        "draining",
        "offline",
        name="p2hs_cloud_shard_status",
        create_type=False,
    )
    project_shard_status_col = ENUM(
        "active",
        "migrating",
        "draining",
        name="p2hs_project_shard_status",
        create_type=False,
    )

    op.create_table(
        "cloud_shards",
        sa.Column("id", PGUUID(as_uuid=True), primary_key=True),
        sa.Column("shard_key", sa.Text(), nullable=False),
        sa.Column("region", sa.Text(), nullable=False, server_default="auto"),
        sa.Column("shard_type", shard_type_col, nullable=False),
        sa.Column(
            "cf_account_id",
            PGUUID(as_uuid=True),
            sa.ForeignKey("cf_accounts.id", ondelete="RESTRICT"),
            nullable=False,
        ),
        sa.Column("d1_database_id", sa.Text(), nullable=True),
        sa.Column("r2_bucket_name", sa.Text(), nullable=True),
        sa.Column("dispatch_namespace", sa.Text(), nullable=True),
        sa.Column(
            "status",
            shard_status_col,
            nullable=False,
            server_default="active",
        ),
        sa.Column("is_accepting", sa.Boolean(), nullable=False, server_default="true"),
        sa.Column("max_projects", sa.Integer(), nullable=False, server_default="10000"),
        sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.func.now()),
        sa.Column("updated_at", sa.DateTime(timezone=True), server_default=sa.func.now()),
        sa.UniqueConstraint("shard_key", name="uq_cloud_shards_shard_key"),
    )
    op.create_index("ix_cloud_shards_shard_type_status", "cloud_shards", ["shard_type", "status"])
    op.create_index("ix_cloud_shards_cf_account_id", "cloud_shards", ["cf_account_id"])

    op.create_table(
        "project_cloud_assignments",
        sa.Column("id", PGUUID(as_uuid=True), primary_key=True),
        sa.Column(
            "project_id",
            PGUUID(as_uuid=True),
            sa.ForeignKey("projects.id", ondelete="CASCADE"),
            nullable=False,
        ),
        sa.Column(
            "shard_id",
            PGUUID(as_uuid=True),
            sa.ForeignKey("cloud_shards.id", ondelete="RESTRICT"),
            nullable=False,
        ),
        sa.Column(
            "status",
            project_shard_status_col,
            nullable=False,
            server_default="active",
        ),
        sa.Column("assigned_at", sa.DateTime(timezone=True), server_default=sa.func.now()),
        sa.UniqueConstraint("project_id", name="uq_project_cloud_assignments_project"),
    )
    op.create_index("ix_project_cloud_assignments_shard_id", "project_cloud_assignments", ["shard_id"])

    # Seed the default shared shard using the first active CF account.
    # If no CF account exists yet, ops should run:
    #   INSERT INTO cloud_shards (id, shard_key, shard_type, cf_account_id, d1_database_id, ...)
    # after inserting the first CF account via POST /api/admin/cf-accounts.
    op.execute("""
        INSERT INTO cloud_shards (
            id, shard_key, region, shard_type, cf_account_id,
            d1_database_id, status, is_accepting, max_projects,
            created_at, updated_at
        )
        SELECT
            '00000000-0000-0000-0000-000000000001',
            'shared-default',
            'auto',
            'shared',
            id,
            'f01a2da8-12b3-494d-85f0-e5c9bf263acb',
            'active',
            true,
            10000,
            NOW(),
            NOW()
        FROM cf_accounts
        WHERE status = 'active'
        ORDER BY created_at
        LIMIT 1
        ON CONFLICT (shard_key) DO NOTHING
    """)


def downgrade() -> None:
    op.drop_table("project_cloud_assignments")
    op.drop_table("cloud_shards")
    op.execute("DROP TYPE IF EXISTS p2hs_project_shard_status")
    op.execute("DROP TYPE IF EXISTS p2hs_cloud_shard_status")
    op.execute("DROP TYPE IF EXISTS p2hs_cloud_shard_type")