"""Add resumable AOI parent and partition operations."""

from alembic import op
import sqlalchemy as sa
from geoalchemy2 import Geometry


revision = "202607260001"
down_revision = "202607160001"
branch_labels = None
depends_on = None


def upgrade() -> None:
    op.create_table(
        "aoi_operations",
        sa.Column("id", sa.UUID(), primary_key=True),
        sa.Column("project_id", sa.UUID(), sa.ForeignKey("projects.id", ondelete="CASCADE"), nullable=False),
        sa.Column("area_id", sa.UUID(), sa.ForeignKey("areas.id", ondelete="SET NULL")),
        sa.Column("parent_job_id", sa.UUID(), sa.ForeignKey("jobs.id", ondelete="SET NULL")),
        sa.Column("operation_type", sa.String(128), nullable=False),
        sa.Column("status", sa.String(32), nullable=False),
        sa.Column("geometry", Geometry("MultiPolygon", srid=4326, spatial_index=False), nullable=False),
        sa.Column("request_json", sa.JSON(), nullable=False),
        sa.Column("plan_json", sa.JSON(), nullable=False),
        sa.Column("result_json", sa.JSON()),
        sa.Column("error_message", sa.Text()),
        sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.func.now()),
        sa.Column("started_at", sa.DateTime(timezone=True)),
        sa.Column("finished_at", sa.DateTime(timezone=True)),
        sa.Column("updated_at", sa.DateTime(timezone=True), server_default=sa.func.now()),
        sa.CheckConstraint("status IN ('queued', 'running', 'partial', 'success', 'failed', 'cancelled')", name="ck_aoi_operations_status"),
    )
    op.create_index("ix_aoi_operations_project_status", "aoi_operations", ["project_id", "status"])
    op.create_index("ix_aoi_operations_geometry", "aoi_operations", ["geometry"], postgresql_using="gist")
    op.create_table(
        "aoi_operation_partitions",
        sa.Column("id", sa.UUID(), primary_key=True),
        sa.Column("operation_id", sa.UUID(), sa.ForeignKey("aoi_operations.id", ondelete="CASCADE"), nullable=False),
        sa.Column("child_job_id", sa.UUID(), sa.ForeignKey("jobs.id", ondelete="SET NULL")),
        sa.Column("partition_key", sa.String(255), nullable=False),
        sa.Column("provider_key", sa.String(120), nullable=False),
        sa.Column("product_key", sa.String(120), nullable=False),
        sa.Column("ordinal", sa.Integer(), nullable=False),
        sa.Column("status", sa.String(32), nullable=False),
        sa.Column("geometry", Geometry("MultiPolygon", srid=4326, spatial_index=False), nullable=False),
        sa.Column("attempt_count", sa.Integer(), nullable=False, server_default="0"),
        sa.Column("max_attempts", sa.Integer(), nullable=False, server_default="3"),
        sa.Column("checkpoint_json", sa.JSON()),
        sa.Column("result_json", sa.JSON()),
        sa.Column("error_message", sa.Text()),
        sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.func.now()),
        sa.Column("started_at", sa.DateTime(timezone=True)),
        sa.Column("finished_at", sa.DateTime(timezone=True)),
        sa.Column("updated_at", sa.DateTime(timezone=True), server_default=sa.func.now()),
        sa.CheckConstraint("status IN ('queued', 'running', 'success', 'failed', 'skipped')", name="ck_aoi_operation_partitions_status"),
        sa.UniqueConstraint("operation_id", "partition_key", name="uq_aoi_operation_partition_key"),
    )
    op.create_index("ix_aoi_operation_partitions_operation_status", "aoi_operation_partitions", ["operation_id", "status"])
    op.create_index("ix_aoi_operation_partitions_geometry", "aoi_operation_partitions", ["geometry"], postgresql_using="gist")


def downgrade() -> None:
    op.drop_table("aoi_operation_partitions")
    op.drop_table("aoi_operations")