"""initial tables

Revision ID: 5c6942374c62
Revises:
Create Date: 2024-02-21 23:56:48.056467

"""

import enum

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

from ai.backend.appproxy.coordinator.models.base import GUID


class FrontendMode(enum.StrEnum):
    WILDCARD_DOMAIN = "wildcard"
    PORT = "port"


class ProxyProtocol(enum.StrEnum):
    HTTP = "http"
    GRPC = "grpc"
    HTTP2 = "h2"
    TCP = "tcp"


class AppMode(enum.StrEnum):
    INTERACTIVE = "interactive"
    INFERENCE = "inference"


appmode_choices = [v.name for v in AppMode]
appmode = postgresql.ENUM(*appmode_choices, name="appmode")

frontendmode_choices = [v.name for v in FrontendMode]
frontendmode = postgresql.ENUM(*frontendmode_choices, name="frontendmode")

proxyprotocol_choices = [v.name for v in ProxyProtocol]
proxyprotocol = postgresql.ENUM(*proxyprotocol_choices, name="proxyprotocol")


# revision identifiers, used by Alembic.
revision = "5c6942374c62"
down_revision = None
branch_labels = None
depends_on = None


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.execute('CREATE EXTENSION IF NOT EXISTS "uuid-ossp";')

    op.create_table(
        "tokens",
        sa.Column("id", GUID(), server_default=sa.text("uuid_generate_v4()"), nullable=False),
        sa.Column("login_session_token", sa.VARCHAR(length=127), nullable=True),
        sa.Column("kernel_host", sa.VARCHAR(length=255), nullable=False),
        sa.Column("kernel_port", sa.INTEGER(), nullable=False),
        sa.Column("session_id", GUID(), nullable=False),
        sa.Column("user_uuid", GUID(), nullable=False),
        sa.Column("group_id", GUID(), nullable=False),
        sa.Column("access_key", sa.TEXT(), nullable=False),
        sa.Column("domain_name", sa.TEXT(), nullable=False),
        sa.PrimaryKeyConstraint("id", name=op.f("pk_tokens")),
    )
    op.create_table(
        "workers",
        sa.Column("id", GUID(), server_default=sa.text("uuid_generate_v4()"), nullable=False),
        sa.Column("authority", sa.String(length=255), nullable=False),
        sa.Column("frontend_mode", frontendmode, nullable=False),
        sa.Column("protocol", proxyprotocol, nullable=False),
        sa.Column("hostname", sa.String(length=1024), nullable=False),
        sa.Column("use_tls", sa.Boolean(), nullable=False),
        sa.Column("api_port", sa.Integer(), nullable=False),
        sa.Column("available_circuits", sa.Integer(), nullable=False),
        sa.Column("occupied_circuits", sa.Integer(), nullable=False),
        sa.Column("port_range", sa.ARRAY(sa.Integer()), nullable=True),
        sa.Column("wildcard_domain", sa.String(length=1024), nullable=True),
        sa.Column("nodes", sa.Integer(), nullable=False),
        sa.Column("accepted_traffics", sa.ARRAY(appmode), nullable=False),
        sa.Column("filtered_apps_only", sa.Boolean(), nullable=False),
        sa.Column(
            "created_at",
            sa.DateTime(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.Column(
            "updated_at",
            sa.DateTime(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.PrimaryKeyConstraint("id", name=op.f("pk_workers")),
        sa.UniqueConstraint("authority", name=op.f("uq_workers_authority")),
    )
    op.create_table(
        "worker_app_filters",
        sa.Column("id", GUID(), server_default=sa.text("uuid_generate_v4()"), nullable=False),
        sa.Column("property_name", sa.VARCHAR(length=96), nullable=False),
        sa.Column("property_value", sa.VARCHAR(length=1024), nullable=False),
        sa.Column("worker", GUID(), server_default=sa.text("uuid_generate_v4()"), nullable=False),
        sa.PrimaryKeyConstraint("id", name=op.f("pk_worker_app_filters")),
        sa.ForeignKeyConstraint(
            ["worker"], ["workers.id"], name=op.f("fk_worker_app_filters_worker_workers")
        ),
        sa.UniqueConstraint(
            "worker",
            "property_name",
            "property_value",
            name="uq_worker_app_filter_worker_property_name_property_value",
        ),
    )
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_table("workers")
    op.drop_table("worker_app_filters")
    op.drop_table("tokens")
    appmode.drop(op.get_bind())
    frontendmode.drop(op.get_bind())
    proxyprotocol.drop(op.get_bind())
    # ### end Alembic commands ###
