"""Update Endpoint and Routing for sokovan

Revision ID: fb01c44db0d6
Revises: 43c7e60c56af
Create Date: 2025-08-29 17:47:24.086038

"""

import sqlalchemy as sa
from alembic import op

# revision identifiers, used by Alembic.
revision = "fb01c44db0d6"
down_revision = "43c7e60c56af"
branch_labels = None
depends_on = None

endpoint_lifecycle_enum_name = "endpointlifecycle"
PENDING = "pending"
SCALING = "scaling"
READY = "ready"
route_status_enum_name = "routestatus"
TERMINATED = "terminated"


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column(
        "endpoints",
        sa.Column("desired_replicas", sa.Integer(), server_default=sa.text("NULL"), nullable=True),
    )
    op.add_column("routings", sa.Column("weight", sa.Integer(), nullable=True))
    op.execute(f"ALTER TYPE {endpoint_lifecycle_enum_name} ADD VALUE '{PENDING}'")
    op.execute(f"ALTER TYPE {endpoint_lifecycle_enum_name} ADD VALUE '{SCALING}'")
    op.execute(f"ALTER TYPE {endpoint_lifecycle_enum_name} ADD VALUE '{READY}'")
    op.execute(f"ALTER TYPE {route_status_enum_name} ADD VALUE '{TERMINATED}'")
    op.create_index(
        "ix_endpoints_unique_name_when_not_destroyed",
        "endpoints",
        ["name", "domain", "project"],
        unique=True,
        postgresql_where=sa.text("lifecycle_stage != 'destroyed'"),
    )
    op.create_index(
        "ix_sessions_unique_name_nonterminal",
        "sessions",
        ["name"],
        unique=True,
        postgresql_where=sa.text("status NOT IN ('ERROR', 'TERMINATED', 'CANCELLED')"),
    )


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_index(
        "ix_sessions_unique_name_nonterminal",
        table_name="sessions",
        postgresql_where=sa.text("status NOT IN ('ERROR', 'TERMINATED', 'CANCELLED')"),
    )
    op.drop_index(
        "ix_endpoints_unique_name_when_not_destroyed",
        table_name="endpoints",
        postgresql_where=sa.text("lifecycle_stage != 'destroyed'"),
    )
    op.drop_column("endpoints", "desired_replicas")
