"""add a new session type INFERENCE

Revision ID: b5be363ab05c
Revises: b6b884fbae1f
Create Date: 2023-02-15 15:24:29.112938

Reference: https://www.postgresql.org/message-id/CANu8FiwwBxZZGX23=Na_7bc4DZ-yzd_poKhaoPmN3+SHG08MAg@mail.gmail.com

"""

import textwrap

from alembic import op
from sqlalchemy.sql import text

from ai.backend.common.types import SessionTypes

# revision identifiers, used by Alembic.
revision = "b5be363ab05c"
down_revision = "b6b884fbae1f"
branch_labels = None
depends_on = None

# select enum_range(null::sessiontypes);
typename = SessionTypes.__name__.lower()


def upgrade() -> None:
    op.execute(text(f"ALTER TYPE {typename} ADD VALUE IF NOT EXISTS 'INFERENCE';"))
    op.execute(
        text(
            textwrap.dedent(
                f"""\
                ALTER TABLE sessions ALTER COLUMN session_type TYPE {typename} USING session_type::text::sessiontypes;
            """
            )
        )
    )


def downgrade() -> None:
    conn = op.get_bind()
    conn.execute(
        text(
            f"UPDATE sessions SET session_type = 'INTERACTIVE'::{typename} WHERE session_type ="
            f" 'INFERENCE'::{typename};"
        )
    )

    cursor = conn.execute(
        text(
            textwrap.dedent(
                """\
                SELECT t.typname, e.enumlabel, e.enumsortorder, e.enumtypid
                FROM pg_type t
                JOIN pg_enum e ON e.enumtypid = t.oid
                WHERE t.typtype = 'e'
                AND e.enumlabel = 'INFERENCE'
                ORDER BY 1, enumsortorder;
            """
            )
        )
    )

    if (row := next(cursor, None)) is None:
        return

    typname, enumlabel, enumsortorder, enumtypid = row

    conn.execute(
        text(
            textwrap.dedent(
                f"""\
                DELETE FROM pg_enum
                WHERE enumtypid = {enumtypid}
                AND enumlabel = 'INFERENCE';
            """
            )
        )
    )
