import sqlalchemy as sa
from alembic import op
from sqlalchemy.dialects.mysql import BIGINT, BINARY, DATETIME

revision = "0005_agent_runtime"
down_revision = "0004_fulfillment"
branch_labels = None
depends_on = None


def identifier(name, primary=False):
    return sa.Column(name, BINARY(16), primary_key=primary, nullable=False)


def upgrade():
    op.create_table(
        "agent_runs",
        identifier("id", True),
        identifier("tenant_id"),
        identifier("actor_id"),
        identifier("store_id"),
        sa.Column("token_hash", BINARY(32), nullable=False),
        sa.Column("request_key", sa.String(64, collation="utf8mb4_bin"), nullable=False),
        sa.Column("request_hash", BINARY(32), nullable=False),
        sa.Column("messages", sa.JSON(), nullable=False),
        sa.Column("status", sa.String(16), nullable=False),
        sa.Column("next_node", sa.String(16), nullable=False),
        sa.Column("graph_version", sa.String(64), nullable=False),
        sa.Column("model_name", sa.String(128), nullable=False),
        sa.Column("model_url", sa.String(512), nullable=False),
        sa.Column("model_credential_env", sa.String(128), nullable=False),
        sa.Column("model_calls", sa.Integer(), nullable=False),
        sa.Column("step", sa.Integer(), nullable=False),
        sa.Column("lease_token", BIGINT(unsigned=True), nullable=False),
        sa.Column("lease_until", DATETIME(fsp=6), nullable=True),
        sa.Column("created_at", DATETIME(fsp=6), nullable=False),
        sa.Column("deadline_at", DATETIME(fsp=6), nullable=False),
        sa.Column("error_code", sa.String(64), nullable=True),
        sa.UniqueConstraint("tenant_id", "id", name="uq_agent_run_tenant_id"),
        sa.UniqueConstraint("tenant_id", "actor_id", "request_key", name="uq_agent_run_request"),
        sa.ForeignKeyConstraint(
            ["tenant_id", "actor_id"], ["principals.tenant_id", "principals.id"]
        ),
        sa.ForeignKeyConstraint(["tenant_id", "store_id"], ["stores.tenant_id", "stores.id"]),
        sa.CheckConstraint(
            "status IN ('QUEUED','RUNNING','COMPLETED','FAILED','CANCELLED')",
            name=op.f("ck_agent_runs_status"),
        ),
    )
    op.create_index("ix_agent_runs_due", "agent_runs", ["status", "created_at", "id"])
    op.create_table(
        "model_attempts",
        identifier("id", True),
        identifier("tenant_id"),
        identifier("run_id"),
        sa.Column("status", sa.String(16), nullable=False),
        sa.Column("usage", sa.JSON(), nullable=True),
        sa.Column("created_at", DATETIME(fsp=6), nullable=False),
        sa.ForeignKeyConstraint(["tenant_id", "run_id"], ["agent_runs.tenant_id", "agent_runs.id"]),
    )


def downgrade():
    op.drop_table("model_attempts")
    op.drop_table("agent_runs")