from __future__ import annotations

from alembic import op
import sqlalchemy as sa

revision = "0002_expand_model_config"
down_revision = "0001_initial_schema"
branch_labels = None
depends_on = None


def upgrade() -> None:
    op.add_column("sys_model", sa.Column("display_name", sa.String(length=100), nullable=True))
    op.add_column(
        "sys_model",
        sa.Column("api_type", sa.String(length=50), nullable=False, server_default="openai_compatible"),
    )
    op.add_column("sys_model", sa.Column("base_url", sa.String(length=255), nullable=True))
    op.add_column(
        "sys_model",
        sa.Column("auth_type", sa.String(length=30), nullable=False, server_default="bearer"),
    )
    op.add_column("sys_model", sa.Column("api_version", sa.String(length=50), nullable=True))
    op.add_column("sys_model", sa.Column("top_p", sa.Numeric(4, 2), nullable=True))
    op.add_column("sys_model", sa.Column("top_k", sa.Integer(), nullable=True))
    op.add_column("sys_model", sa.Column("presence_penalty", sa.Numeric(4, 2), nullable=True))
    op.add_column("sys_model", sa.Column("frequency_penalty", sa.Numeric(4, 2), nullable=True))
    op.add_column("sys_model", sa.Column("context_window", sa.Integer(), nullable=True))
    op.add_column(
        "sys_model",
        sa.Column("stream_enabled", sa.Integer(), nullable=False, server_default="1"),
    )
    op.add_column("sys_model", sa.Column("response_format", sa.String(length=50), nullable=True))
    op.add_column("sys_model", sa.Column("extra_params", sa.Text(), nullable=True))
    op.add_column("sys_model", sa.Column("remark", sa.String(length=255), nullable=True))

    op.alter_column("sys_model", "api_type", server_default=None)
    op.alter_column("sys_model", "auth_type", server_default=None)
    op.alter_column("sys_model", "stream_enabled", server_default=None)


def downgrade() -> None:
    op.drop_column("sys_model", "remark")
    op.drop_column("sys_model", "extra_params")
    op.drop_column("sys_model", "response_format")
    op.drop_column("sys_model", "stream_enabled")
    op.drop_column("sys_model", "context_window")
    op.drop_column("sys_model", "frequency_penalty")
    op.drop_column("sys_model", "presence_penalty")
    op.drop_column("sys_model", "top_k")
    op.drop_column("sys_model", "top_p")
    op.drop_column("sys_model", "api_version")
    op.drop_column("sys_model", "auth_type")
    op.drop_column("sys_model", "base_url")
    op.drop_column("sys_model", "api_type")
    op.drop_column("sys_model", "display_name")
