"""2.17.0

Revision ID: 83bc2ad3ffcc
Revises: 981d588fe248
Create Date: 2025-10-30 14:16:53.639006

"""

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

# revision identifiers, used by Alembic.
revision = "83bc2ad3ffcc"
down_revision = "981d588fe248"
branch_labels = None
depends_on = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        "resource",
        sa.Column("id", sa.Integer(), nullable=False),
        sa.Column("type", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("name", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column(
            "timestamp_created", sa.DateTime(timezone=True), nullable=False
        ),
        sa.Column("host", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
        sa.Column(
            "jobs_local_dir",
            sqlmodel.sql.sqltypes.AutoString(),
            nullable=False,
        ),
        sa.Column(
            "jobs_runner_config",
            postgresql.JSONB(astext_type=sa.Text()),
            server_default="{}",
            nullable=False,
        ),
        sa.Column(
            "jobs_slurm_python_worker",
            sqlmodel.sql.sqltypes.AutoString(),
            nullable=True,
        ),
        sa.Column("jobs_poll_interval", sa.Integer(), nullable=False),
        sa.Column(
            "tasks_local_dir",
            sqlmodel.sql.sqltypes.AutoString(),
            nullable=False,
        ),
        sa.Column(
            "tasks_python_config",
            postgresql.JSONB(astext_type=sa.Text()),
            server_default="{}",
            nullable=False,
        ),
        sa.Column(
            "tasks_pixi_config",
            postgresql.JSONB(astext_type=sa.Text()),
            server_default="{}",
            nullable=False,
        ),
        sa.CheckConstraint(
            "(type = 'local') OR (jobs_slurm_python_worker IS NOT NULL)",
            name=op.f("ck_resource_jobs_slurm_python_worker_set"),
        ),
        sa.CheckConstraint(
            "type IN ('local', 'slurm_sudo', 'slurm_ssh')",
            name=op.f("ck_resource_correct_type"),
        ),
        sa.PrimaryKeyConstraint("id", name=op.f("pk_resource")),
        sa.UniqueConstraint("name", name=op.f("uq_resource_name")),
    )
    op.create_table(
        "profile",
        sa.Column("id", sa.Integer(), nullable=False),
        sa.Column("resource_id", sa.Integer(), nullable=False),
        sa.Column(
            "resource_type", sqlmodel.sql.sqltypes.AutoString(), nullable=False
        ),
        sa.Column("name", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column(
            "username", sqlmodel.sql.sqltypes.AutoString(), nullable=True
        ),
        sa.Column(
            "ssh_key_path", sqlmodel.sql.sqltypes.AutoString(), nullable=True
        ),
        sa.Column(
            "jobs_remote_dir",
            sqlmodel.sql.sqltypes.AutoString(),
            nullable=True,
        ),
        sa.Column(
            "tasks_remote_dir",
            sqlmodel.sql.sqltypes.AutoString(),
            nullable=True,
        ),
        sa.ForeignKeyConstraint(
            ["resource_id"],
            ["resource.id"],
            name=op.f("fk_profile_resource_id_resource"),
            ondelete="RESTRICT",
        ),
        sa.PrimaryKeyConstraint("id", name=op.f("pk_profile")),
        sa.UniqueConstraint("name", name=op.f("uq_profile_name")),
    )
    with op.batch_alter_table("projectv2", schema=None) as batch_op:
        batch_op.add_column(
            sa.Column("resource_id", sa.Integer(), nullable=True)
        )
        batch_op.create_foreign_key(
            batch_op.f("fk_projectv2_resource_id_resource"),
            "resource",
            ["resource_id"],
            ["id"],
            ondelete="RESTRICT",
        )

    with op.batch_alter_table("taskgroupv2", schema=None) as batch_op:
        batch_op.add_column(
            sa.Column("resource_id", sa.Integer(), nullable=True)
        )
        batch_op.create_foreign_key(
            batch_op.f("fk_taskgroupv2_resource_id_resource"),
            "resource",
            ["resource_id"],
            ["id"],
            ondelete="RESTRICT",
        )

    with op.batch_alter_table("user_oauth", schema=None) as batch_op:
        batch_op.add_column(
            sa.Column("profile_id", sa.Integer(), nullable=True)
        )
        batch_op.add_column(
            sa.Column(
                "project_dir",
                sa.String(),
                server_default="/PLACEHOLDER",
                nullable=False,
            )
        )
        batch_op.add_column(
            sa.Column(
                "slurm_accounts",
                postgresql.ARRAY(sa.String()),
                server_default="{}",
                nullable=True,
            )
        )
        batch_op.create_foreign_key(
            batch_op.f("fk_user_oauth_profile_id_profile"),
            "profile",
            ["profile_id"],
            ["id"],
            ondelete="RESTRICT",
        )
        batch_op.drop_column("username")

    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table("user_oauth", schema=None) as batch_op:
        batch_op.add_column(
            sa.Column(
                "username", sa.VARCHAR(), autoincrement=False, nullable=True
            )
        )
        batch_op.drop_constraint(
            batch_op.f("fk_user_oauth_profile_id_profile"), type_="foreignkey"
        )
        batch_op.drop_column("slurm_accounts")
        batch_op.drop_column("project_dir")
        batch_op.drop_column("profile_id")

    with op.batch_alter_table("taskgroupv2", schema=None) as batch_op:
        batch_op.drop_constraint(
            batch_op.f("fk_taskgroupv2_resource_id_resource"),
            type_="foreignkey",
        )
        batch_op.drop_column("resource_id")

    with op.batch_alter_table("projectv2", schema=None) as batch_op:
        batch_op.drop_constraint(
            batch_op.f("fk_projectv2_resource_id_resource"), type_="foreignkey"
        )
        batch_op.drop_column("resource_id")

    op.drop_table("profile")
    op.drop_table("resource")
    # ### end Alembic commands ###
