"""add-keypair-resource-policy

Revision ID: dbc1e053b880
Revises: 2b0931e4a059
Create Date: 2019-02-07 15:30:54.861821

"""

import sqlalchemy as sa
from alembic import op
from sqlalchemy.dialects import postgresql
from sqlalchemy.sql import text

from ai.backend.common.types import DefaultForUnspecified

# revision identifiers, used by Alembic.
revision = "dbc1e053b880"
down_revision = "2b0931e4a059"
branch_labels = None
depends_on = None


default_for_unspecified_choices = [*map(str, DefaultForUnspecified)]
default_for_unspecified = postgresql.ENUM(
    *default_for_unspecified_choices,
    name="default_for_unspecified",
)


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        "keypair_resource_policies",
        sa.Column("name", sa.String(length=256), nullable=False),
        sa.Column(
            "created_at", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=True
        ),
        sa.Column(
            "default_for_unspecified",
            type_=sa.Enum(*default_for_unspecified_choices, name="default_for_unspecified"),
            nullable=False,
        ),
        sa.Column("total_resource_slots", postgresql.JSONB(astext_type=sa.Text()), nullable=False),
        sa.Column("max_concurrent_sessions", sa.Integer(), nullable=False),
        sa.Column("max_containers_per_session", sa.Integer(), nullable=False),
        sa.Column("max_vfolder_count", sa.Integer(), nullable=False),
        sa.Column("max_vfolder_size", sa.BigInteger(), nullable=False),
        sa.Column("allowed_vfolder_hosts", postgresql.ARRAY(sa.String()), nullable=False),
        sa.PrimaryKeyConstraint("name", name=op.f("pk_keypair_resource_policies")),
    )
    # Create a default resource policy
    query = """
    INSERT INTO keypair_resource_policies
    VALUES (
        'default',
        now(),
        'UNLIMITED',
        '{}'::jsonb,
        30,
        1,
        10,
        0,
        array['local']
    );
    """
    connection = op.get_bind()
    connection.execute(text(query))
    query = """
    UPDATE keypairs
    SET resource_policy = 'default';
    """
    connection.execute(text(query))
    print("\n!!! NOTICE !!!\n")
    print("Created a default resource policy and linked all keypairs to it.")
    print("Please inspect and adjust it!\n")
    op.alter_column(
        "keypairs",
        "resource_policy",
        existing_type=sa.VARCHAR(),
        type_=sa.String(length=256),
        nullable=False,
    )
    op.create_foreign_key(
        op.f("fk_keypairs_resource_policy_keypair_resource_policies"),
        "keypairs",
        "keypair_resource_policies",
        ["resource_policy"],
        ["name"],
    )
    op.drop_column("keypairs", "concurrency_limit")
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_constraint(
        op.f("fk_keypairs_resource_policy_keypair_resource_policies"),
        "keypairs",
        type_="foreignkey",
    )
    op.add_column(
        "keypairs", sa.Column("concurrency_limit", sa.INTEGER(), autoincrement=False, nullable=True)
    )
    op.alter_column(
        "keypairs",
        "resource_policy",
        existing_type=sa.String(length=256),
        type_=sa.VARCHAR(),
        nullable=True,
    )
    op.drop_table("keypair_resource_policies")
    default_for_unspecified.drop(op.get_bind())
    # ### end Alembic commands ###
