"""make workspace_id non-nullable for workspace-scoped models

Revision ID: 0fd1f09cd98b
Revises: b23b877ee161
Create Date: 2025-12-04 14:14:05.910579

"""

from collections.abc import Sequence

import sqlalchemy as sa

from alembic import op

# revision identifiers, used by Alembic.
revision: str = "0fd1f09cd98b"
down_revision: str | None = "b23b877ee161"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None

# Tables that should not contain NULL workspace_id values before enforcing the
# NOT NULL constraint.
WORKSPACE_SCOPED_TABLES: tuple[str, ...] = (
    "agent_preset",
    "case_duration",
    "case_duration_definition",
    "case_tag",
    "secret",
    "tag",
    "workflow",
    "workflow_folder",
)


def upgrade() -> None:
    _assert_no_null_workspace_ids()

    # ### commands auto generated by Alembic - please adjust! ###
    op.alter_column(
        "agent_preset", "workspace_id", existing_type=sa.UUID(), nullable=False
    )
    op.alter_column(
        "case_duration", "workspace_id", existing_type=sa.UUID(), nullable=False
    )
    op.alter_column(
        "case_duration_definition",
        "workspace_id",
        existing_type=sa.UUID(),
        nullable=False,
    )
    op.alter_column("case_tag", "workspace_id", existing_type=sa.UUID(), nullable=False)
    op.alter_column("secret", "workspace_id", existing_type=sa.UUID(), nullable=False)
    op.alter_column("tag", "workspace_id", existing_type=sa.UUID(), nullable=False)
    op.alter_column("workflow", "workspace_id", existing_type=sa.UUID(), nullable=False)
    op.alter_column(
        "workflow_folder", "workspace_id", existing_type=sa.UUID(), nullable=False
    )
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.alter_column(
        "workflow_folder", "workspace_id", existing_type=sa.UUID(), nullable=True
    )
    op.alter_column("workflow", "workspace_id", existing_type=sa.UUID(), nullable=True)
    op.alter_column("tag", "workspace_id", existing_type=sa.UUID(), nullable=True)
    op.alter_column("secret", "workspace_id", existing_type=sa.UUID(), nullable=True)
    op.alter_column("case_tag", "workspace_id", existing_type=sa.UUID(), nullable=True)
    op.alter_column(
        "case_duration_definition",
        "workspace_id",
        existing_type=sa.UUID(),
        nullable=True,
    )
    op.alter_column(
        "case_duration", "workspace_id", existing_type=sa.UUID(), nullable=True
    )
    op.alter_column(
        "agent_preset", "workspace_id", existing_type=sa.UUID(), nullable=True
    )
    # ### end Alembic commands ###


def _assert_no_null_workspace_ids() -> None:
    """Ensure workspace_id is populated before enforcing NOT NULL.

    Earlier migrations created these columns as nullable, so deployments with
    legacy data might still have rows with NULL workspace_id values. Tightening
    the constraint without backfilling would fail mid-migration; this guard
    fails fast with a clear error message so operators can remediate their
    data before rerunning the migration.
    """

    conn = op.get_bind()

    tables_with_nulls: list[tuple[str, int]] = []
    for table in WORKSPACE_SCOPED_TABLES:
        count = conn.execute(
            sa.text(f"SELECT COUNT(*) FROM {table} WHERE workspace_id IS NULL")
        ).scalar_one()
        if count:
            tables_with_nulls.append((table, count))

    if tables_with_nulls:
        details = ", ".join(
            f"{table} ({count} rows)" for table, count in tables_with_nulls
        )
        raise RuntimeError(
            "Found NULL workspace_id values; backfill before rerunning migration: "
            f"{details}"
        )
