"""add kubernetes runner: remote_task_results table, kubernetes_config column on jobs and registered_jobs, drop runner_mode check constraints

Revision ID: 1798e1da266d
Revises: ce57a9d1f1a9
Create Date: 2026-06-16 02:54:59.472212

"""

from collections.abc import Sequence

import sqlalchemy as sa
import sqlmodel
from alembic import op

# revision identifiers, used by Alembic.
revision: str = "1798e1da266d"
down_revision: str | Sequence[str] | None = "76dc173ce459"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None


def upgrade() -> None:
    """Upgrade schema."""
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        "remote_task_results",
        sa.Column("task_id", sa.BigInteger(), nullable=False),
        sa.Column("run_epoch", sa.BigInteger(), nullable=False),
        sa.Column("success", sa.Boolean(), nullable=False),
        sa.Column("result_ref", sa.JSON(), nullable=True),
        sa.Column("log_path", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
        sa.Column("error", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
        sa.Column("created_at", sa.DateTime(), nullable=False),
        sa.ForeignKeyConstraint(
            ["task_id"],
            ["tasks.id"],
        ),
        sa.PrimaryKeyConstraint("task_id", "run_epoch"),
    )
    op.add_column("jobs", sa.Column("kubernetes_config", sa.JSON(), nullable=True))
    op.add_column("registered_jobs", sa.Column("kubernetes_config", sa.JSON(), nullable=True))
    # ### end Alembic commands ###

    # Alembic autogenerate cannot detect CHECK-constraint changes. Drop the
    # runner_mode CHECKs instead of widening them: runner_mode is validated by
    # the RunnerMode Literal + the CLI's choices=, so adding a runner mode
    # (e.g. 'kubernetes') needs no further migration.
    with op.batch_alter_table("jobs") as batch:
        batch.drop_constraint("ck_jobs_runner_mode", type_="check")
    with op.batch_alter_table("registered_jobs") as batch:
        batch.drop_constraint("ck_registered_jobs_runner_mode", type_="check")


def downgrade() -> None:
    """Downgrade schema."""
    # Recreate the runner_mode CHECK constraints as they were before this
    # migration (subprocess + docker only).
    with op.batch_alter_table("registered_jobs") as batch:
        batch.create_check_constraint("ck_registered_jobs_runner_mode", "runner_mode IN ('subprocess', 'docker')")
    with op.batch_alter_table("jobs") as batch:
        batch.create_check_constraint("ck_jobs_runner_mode", "runner_mode IN ('subprocess', 'docker')")

    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_column("registered_jobs", "kubernetes_config")
    op.drop_column("jobs", "kubernetes_config")
    op.drop_table("remote_task_results")
    # ### end Alembic commands ###
