"""add-resource-presets

Revision ID: 8e660aa31fe3
Revises: 01456c812164
Create Date: 2019-03-30 01:45:07.525096

"""

from decimal import Decimal
from typing import Any, cast

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

from ai.backend.common.types import BinarySize, ResourceSlot
from ai.backend.manager.models.base import ResourceSlotColumn
from ai.backend.manager.models.resource_policy import keypair_resource_policies

# revision identifiers, used by Alembic.
revision = "8e660aa31fe3"
down_revision = "01456c812164"
branch_labels = None
depends_on = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        "resource_presets",
        sa.Column("name", sa.String(length=256), nullable=False),
        sa.Column("resource_slots", ResourceSlotColumn(), nullable=False),
        sa.PrimaryKeyConstraint("name", name=op.f("pk_resource_presets")),
    )
    # Add initial fixtures for resource presets
    queries = """
    INSERT INTO resource_presets
    VALUES (
        'small',
        '{"cpu":"1","mem":"2147483648"}'::jsonb
    );
    INSERT INTO resource_presets
    VALUES (
        'small-gpu',
        '{"cpu":"1","mem":"2147483648","cuda.device":"1","cuda.shares":"0.5"}'::jsonb
    );
    INSERT INTO resource_presets
    VALUES (
        'medium',
        '{"cpu":"2","mem":"4294967296"}'::jsonb
    );
    INSERT INTO resource_presets
    VALUES (
        'medium-gpu',
        '{"cpu":"2","mem":"4294967296","cuda.device":"1","cuda.shares":"1.0"}'::jsonb
    );
    INSERT INTO resource_presets
    VALUES (
        'large',
        '{"cpu":"4","mem":"8589934592"}'::jsonb
    );
    INSERT INTO resource_presets
    VALUES (
        'large-gpu',
        '{"cpu":"4","mem":"8589934592","cuda.device":"2","cuda.shares":"2.0"}'::jsonb
    );
    """
    connection = op.get_bind()
    for query in queries.split(";"):
        if len(query.strip()) == 0:
            continue
        connection.execute(text(query))

    query = """
    SELECT name, total_resource_slots
    FROM keypair_resource_policies
    """
    result = connection.execute(text(query))
    updates = []
    for row in result:
        row_mapping = cast(dict[str, Any], row._mapping)
        converted = ResourceSlot(row_mapping["total_resource_slots"])
        if "mem" in converted:
            converted["mem"] = Decimal(BinarySize.from_str(converted["mem"]))
            updates.append((
                row_mapping["name"],
                converted,
            ))
    for name, slots in updates:
        update_query = (
            sa.update(keypair_resource_policies)
            .values(total_resource_slots=slots)
            .where(keypair_resource_policies.c.name == name)
        )
        connection.execute(update_query)

    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_table("resource_presets")
    # ### end Alembic commands ###
