"""add gpu_lipsync_tasks and gpu_workers tables for MuseTalk reverse-poll worker

Revision ID: 081_add_gpu_lipsync
Revises: 080_edit_plan_clips_atom_clip_id
Create Date: 2026-09-18
"""

import sqlalchemy as sa

from alembic import op

revision = "081_add_gpu_lipsync"
down_revision = "080_edit_plan_clips_atom_clip_id"
branch_labels = None
depends_on = None


def upgrade() -> None:
    # GPU Worker 注册表
    op.create_table(
        "gpu_workers",
        sa.Column("worker_id", sa.String(100), primary_key=True),
        sa.Column("hostname", sa.String(200), nullable=False, server_default=""),
        sa.Column("gpu_name", sa.String(200), nullable=False, server_default=""),
        sa.Column("free_vram_mb", sa.Integer(), nullable=False, server_default=sa.text("0")),
        sa.Column("capabilities", sa.String(500), nullable=False, server_default=""),
        sa.Column("last_heartbeat_at", sa.DateTime(), nullable=True, index=True),
        sa.Column("created_at", sa.DateTime(), nullable=False, server_default=sa.func.now()),
    )

    # GPU 口型同步任务表
    op.create_table(
        "gpu_lipsync_tasks",
        sa.Column("id", sa.String(36), primary_key=True),
        sa.Column("lipsync_job_id", sa.String(36), nullable=False, server_default="", index=True),
        sa.Column("user_id", sa.String(36), nullable=False, server_default="", index=True),
        sa.Column("project_id", sa.String(36), nullable=False, server_default="", index=True),
        sa.Column("video_url", sa.Text(), nullable=False),
        sa.Column("audio_url", sa.Text(), nullable=False),
        sa.Column("result_url", sa.Text(), nullable=False, server_default=""),
        sa.Column("result_duration", sa.Float(), nullable=False, server_default=sa.text("0.0")),
        sa.Column("status", sa.String(20), nullable=False, server_default="pending", index=True),
        sa.Column("worker_id", sa.String(100), nullable=False, server_default="", index=True),
        sa.Column("attempt", sa.Integer(), nullable=False, server_default=sa.text("0")),
        sa.Column("error_msg", sa.Text(), nullable=False, server_default=""),
        sa.Column("created_at", sa.DateTime(), nullable=False, server_default=sa.func.now()),
        sa.Column("started_at", sa.DateTime(), nullable=True),
        sa.Column("finished_at", sa.DateTime(), nullable=True),
        sa.Column("updated_at", sa.DateTime(), nullable=False, server_default=sa.func.now()),
        sa.Column("last_heartbeat_at", sa.DateTime(), nullable=True),
    )
    op.create_index("ix_gpu_lipsync_status_created", "gpu_lipsync_tasks", ["status", "created_at"])


def downgrade() -> None:
    op.drop_index("ix_gpu_lipsync_status_created", table_name="gpu_lipsync_tasks")
    op.drop_table("gpu_lipsync_tasks")
    op.drop_table("gpu_workers")