"""initial schema

Revision ID: b0839badf97f
Revises:
Create Date: 2026-04-29 18:46:22.466475

"""

from collections.abc import Sequence

import sqlalchemy as sa
import sqlmodel
from alembic import op

# revision identifiers, used by Alembic.
revision: str = "b0839badf97f"
down_revision: str | Sequence[str] | None = None
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None


def upgrade() -> None:
    """Upgrade schema."""
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        "dependencies",
        sa.Column("previous_id", sa.BigInteger(), nullable=False),
        sa.Column("previous_type", sa.String(), nullable=False),
        sa.Column("next_id", sa.BigInteger(), nullable=False),
        sa.Column("next_type", sa.String(), nullable=False),
        sa.Column("created_at", sa.DateTime(), nullable=False),
        sa.PrimaryKeyConstraint("previous_id", "previous_type", "next_id", "next_type"),
    )
    op.create_index(op.f("ix_dependencies_next_id"), "dependencies", ["next_id"], unique=False)
    op.create_index(op.f("ix_dependencies_previous_id"), "dependencies", ["previous_id"], unique=False)
    op.create_table(
        "registered_jobs",
        sa.Column("id", sa.BigInteger(), nullable=False),
        sa.Column("name", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("entrypoint", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("enabled", sa.Boolean(), server_default="1", nullable=False),
        sa.Column("schedule", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
        sa.Column("default_kwargs", sa.JSON(), nullable=True),
        sa.Column("preservation_mode", sa.String(), nullable=True),
        sa.CheckConstraint(
            "preservation_mode IN ('NONE', 'FULL')",
            name="ck_registered_jobs_preservation_mode",
        ),
        sa.Column("next_run_at", sa.DateTime(), nullable=True),
        sa.Column("created_at", sa.DateTime(), nullable=False),
        sa.Column("updated_at", sa.DateTime(), nullable=False),
        sa.PrimaryKeyConstraint("id"),
        sa.UniqueConstraint("name"),
    )
    op.create_index(op.f("ix_registered_jobs_name"), "registered_jobs", ["name"], unique=False)
    op.create_index(op.f("ix_registered_jobs_next_run_at"), "registered_jobs", ["next_run_at"], unique=False)
    op.create_table(
        "table_context_refs",
        sa.Column("table_name", sa.String(), nullable=False),
        sa.Column("context_id", sa.BigInteger(), nullable=False),
        sa.Column("advisory_id", sa.BigInteger(), nullable=False),
        sa.PrimaryKeyConstraint("table_name", "context_id"),
    )
    op.create_table(
        "table_pin_refs",
        sa.Column("table_name", sa.String(), nullable=False),
        sa.Column("task_id", sa.BigInteger(), nullable=False),
        sa.PrimaryKeyConstraint("table_name", "task_id"),
    )
    op.create_table(
        "table_registry",
        sa.Column("table_name", sa.String(), nullable=False),
        sa.Column("job_id", sa.BigInteger(), nullable=True),
        sa.Column("task_id", sa.BigInteger(), nullable=True),
        sa.Column("run_id", sa.BigInteger(), nullable=True),
        sa.Column("created_at", sa.DateTime(), nullable=False),
        sa.Column("schema_doc", sa.Text(), nullable=True),
        sa.PrimaryKeyConstraint("table_name"),
    )
    op.create_index(op.f("ix_table_registry_created_at"), "table_registry", ["created_at"], unique=False)
    op.create_index(op.f("ix_table_registry_job_id"), "table_registry", ["job_id"], unique=False)
    op.create_table(
        "table_run_refs",
        sa.Column("table_name", sa.String(), nullable=False),
        sa.Column("run_id", sa.String(), nullable=False),
        sa.PrimaryKeyConstraint("table_name", "run_id"),
    )
    op.create_table(
        "workers",
        sa.Column("id", sa.BigInteger(), nullable=False),
        sa.Column("hostname", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("pid", sa.Integer(), nullable=False),
        sa.Column("status", sa.String(), nullable=False),
        sa.Column("created_at", sa.DateTime(), nullable=False),
        sa.Column("started_at", sa.DateTime(), nullable=False),
        sa.Column("last_heartbeat", sa.DateTime(), nullable=False),
        sa.Column("tasks_completed", sa.Integer(), nullable=False),
        sa.Column("tasks_failed", sa.Integer(), nullable=False),
        sa.CheckConstraint(
            "status IN ('ACTIVE', 'IDLE', 'STOPPING', 'STOPPED')",
            name="ck_workers_status",
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index(op.f("ix_workers_hostname"), "workers", ["hostname"], unique=False)
    op.create_index(op.f("ix_workers_last_heartbeat"), "workers", ["last_heartbeat"], unique=False)
    op.create_index(op.f("ix_workers_status"), "workers", ["status"], unique=False)
    op.create_table(
        "jobs",
        sa.Column("id", sa.BigInteger(), nullable=False),
        sa.Column("name", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("status", sa.String(), nullable=False),
        sa.Column("run_type", sa.String(), nullable=False),
        sa.Column("registered_job_id", sa.BigInteger(), nullable=True),
        sa.Column("preservation_mode", sa.String(), server_default="NONE", nullable=False),
        sa.Column("created_at", sa.DateTime(), nullable=False),
        sa.Column("started_at", sa.DateTime(), nullable=True),
        sa.Column("completed_at", sa.DateTime(), nullable=True),
        sa.Column("error", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
        sa.CheckConstraint(
            "status IN ('PENDING', 'RUNNING', 'COMPLETED', 'FAILED', 'CANCELLED')",
            name="ck_jobs_status",
        ),
        sa.CheckConstraint("run_type IN ('SCHEDULED', 'MANUAL')", name="ck_jobs_run_type"),
        sa.CheckConstraint(
            "preservation_mode IN ('NONE', 'FULL')",
            name="ck_jobs_preservation_mode",
        ),
        sa.ForeignKeyConstraint(
            ["registered_job_id"],
            ["registered_jobs.id"],
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index(op.f("ix_jobs_created_at"), "jobs", ["created_at"], unique=False)
    op.create_index(op.f("ix_jobs_name"), "jobs", ["name"], unique=False)
    op.create_index(op.f("ix_jobs_registered_job_id"), "jobs", ["registered_job_id"], unique=False)
    op.create_index(op.f("ix_jobs_status"), "jobs", ["status"], unique=False)
    op.create_table(
        "groups",
        sa.Column("id", sa.BigInteger(), nullable=False),
        sa.Column("job_id", sa.BigInteger(), nullable=True),
        sa.Column("parent_group_id", sa.BigInteger(), nullable=True),
        sa.Column("name", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("created_at", sa.DateTime(), nullable=False),
        sa.ForeignKeyConstraint(
            ["job_id"],
            ["jobs.id"],
        ),
        sa.ForeignKeyConstraint(
            ["parent_group_id"],
            ["groups.id"],
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index(op.f("ix_groups_job_id"), "groups", ["job_id"], unique=False)
    op.create_index(op.f("ix_groups_parent_group_id"), "groups", ["parent_group_id"], unique=False)
    op.create_table(
        "tasks",
        sa.Column("id", sa.BigInteger(), nullable=False),
        sa.Column("job_id", sa.BigInteger(), nullable=True),
        sa.Column("group_id", sa.BigInteger(), nullable=True),
        sa.Column("entrypoint", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("name", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("kwargs", sa.JSON(), nullable=True),
        sa.Column("status", sa.String(), nullable=False),
        sa.Column("created_at", sa.DateTime(), nullable=False),
        sa.Column("claimed_at", sa.DateTime(), nullable=True),
        sa.Column("started_at", sa.DateTime(), nullable=True),
        sa.Column("completed_at", sa.DateTime(), nullable=True),
        sa.Column("worker_id", sa.BigInteger(), nullable=True),
        sa.Column("result", sa.JSON(), nullable=True),
        sa.Column("log_path", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
        sa.Column("error", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
        sa.Column("max_retries", sa.Integer(), nullable=False),
        sa.Column("attempt", sa.Integer(), nullable=False),
        sa.Column("retry_after", sa.DateTime(), nullable=True),
        sa.Column("run_ids", sa.JSON(), server_default="[]", nullable=False),
        sa.Column("run_statuses", sa.JSON(), server_default="[]", nullable=False),
        sa.CheckConstraint(
            "status IN ('PENDING', 'CLAIMED', 'RUNNING', 'COMPLETED', 'FAILED', 'CANCELLED', 'PENDING_CLEANUP')",
            name="ck_tasks_status",
        ),
        sa.ForeignKeyConstraint(
            ["group_id"],
            ["groups.id"],
        ),
        sa.ForeignKeyConstraint(
            ["job_id"],
            ["jobs.id"],
        ),
        sa.ForeignKeyConstraint(
            ["worker_id"],
            ["workers.id"],
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index(op.f("ix_tasks_created_at"), "tasks", ["created_at"], unique=False)
    op.create_index(op.f("ix_tasks_group_id"), "tasks", ["group_id"], unique=False)
    op.create_index(op.f("ix_tasks_job_id"), "tasks", ["job_id"], unique=False)
    op.create_index(op.f("ix_tasks_status"), "tasks", ["status"], unique=False)
    op.create_index(op.f("ix_tasks_worker_id"), "tasks", ["worker_id"], unique=False)
    # ### end Alembic commands ###


def downgrade() -> None:
    """Downgrade schema."""
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_index(op.f("ix_tasks_worker_id"), table_name="tasks")
    op.drop_index(op.f("ix_tasks_status"), table_name="tasks")
    op.drop_index(op.f("ix_tasks_job_id"), table_name="tasks")
    op.drop_index(op.f("ix_tasks_group_id"), table_name="tasks")
    op.drop_index(op.f("ix_tasks_created_at"), table_name="tasks")
    op.drop_table("tasks")
    op.drop_index(op.f("ix_groups_parent_group_id"), table_name="groups")
    op.drop_index(op.f("ix_groups_job_id"), table_name="groups")
    op.drop_table("groups")
    op.drop_index(op.f("ix_jobs_status"), table_name="jobs")
    op.drop_index(op.f("ix_jobs_registered_job_id"), table_name="jobs")
    op.drop_index(op.f("ix_jobs_name"), table_name="jobs")
    op.drop_index(op.f("ix_jobs_created_at"), table_name="jobs")
    op.drop_table("jobs")
    op.drop_index(op.f("ix_workers_status"), table_name="workers")
    op.drop_index(op.f("ix_workers_last_heartbeat"), table_name="workers")
    op.drop_index(op.f("ix_workers_hostname"), table_name="workers")
    op.drop_table("workers")
    op.drop_table("table_run_refs")
    op.drop_index(op.f("ix_table_registry_job_id"), table_name="table_registry")
    op.drop_index(op.f("ix_table_registry_created_at"), table_name="table_registry")
    op.drop_table("table_registry")
    op.drop_table("table_pin_refs")
    op.drop_table("table_context_refs")
    op.drop_index(op.f("ix_registered_jobs_next_run_at"), table_name="registered_jobs")
    op.drop_index(op.f("ix_registered_jobs_name"), table_name="registered_jobs")
    op.drop_table("registered_jobs")
    op.drop_index(op.f("ix_dependencies_previous_id"), table_name="dependencies")
    op.drop_index(op.f("ix_dependencies_next_id"), table_name="dependencies")
    op.drop_table("dependencies")
    # ### end Alembic commands ###
