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

revision = "0022_run_inputs"
down_revision = "0021_billing_statements"
branch_labels = None
depends_on = None


def upgrade():
    op.add_column("agent_runs", sa.Column("version", mysql.BIGINT(unsigned=True),
                                         nullable=False, server_default="1"))
    op.alter_column("agent_runs", "version", server_default=None,
                    existing_type=mysql.BIGINT(unsigned=True))
    op.add_column("agent_runs", sa.Column("waiting", sa.JSON(), nullable=True))
    op.drop_constraint(op.f("ck_agent_runs_status"), "agent_runs", type_="check")
    op.create_check_constraint("status", "agent_runs",
        "status IN ('QUEUED','RUNNING','WAITING_USER','COMPLETED','FAILED','CANCELLED')")
    # 每次等待只能接受一份有效输入，输入收据与运行恢复使用同一事务。
    op.create_table(
        "run_input_events",
        sa.Column("id", mysql.BINARY(16), primary_key=True),
        sa.Column("tenant_id", mysql.BINARY(16), nullable=False),
        sa.Column("run_id", mysql.BINARY(16), nullable=False),
        sa.Column("actor_id", mysql.BINARY(16), nullable=False),
        sa.Column("interrupt_id", mysql.BINARY(16), nullable=False),
        sa.Column("request_key", sa.String(64, collation="utf8mb4_bin"), nullable=False),
        sa.Column("request_hash", mysql.BINARY(32), nullable=False),
        sa.Column("payload", sa.JSON(), nullable=False),
        sa.Column("result_version", mysql.BIGINT(unsigned=True), nullable=False),
        sa.Column("created_at", mysql.DATETIME(fsp=6), nullable=False),
        sa.ForeignKeyConstraint(["tenant_id", "run_id"], ["agent_runs.tenant_id", "agent_runs.id"]),
        sa.UniqueConstraint("tenant_id", "run_id", "request_key", name="uq_run_input_request"),
        sa.UniqueConstraint("tenant_id", "run_id", "interrupt_id", name="uq_run_input_wait"),
    )


def downgrade():
    op.drop_table("run_input_events")
    op.drop_constraint(op.f("ck_agent_runs_status"), "agent_runs", type_="check")
    op.create_check_constraint("status", "agent_runs",
        "status IN ('QUEUED','RUNNING','COMPLETED','FAILED','CANCELLED')")
    op.drop_column("agent_runs", "waiting")
    op.drop_column("agent_runs", "version")