"""agent_sessions

Revision ID: f856d9b7ee5a
Revises: 47e44115516c
Create Date: 2026-01-15 18:28:01.938505

"""

from collections.abc import Sequence

import sqlalchemy as sa
from sqlalchemy.dialects import postgresql

from alembic import op

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


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        "agent_session",
        sa.Column("id", sa.UUID(), nullable=False),
        sa.Column("title", sa.String(length=200), nullable=False),
        sa.Column("created_by", sa.UUID(), nullable=True),
        sa.Column("entity_type", sa.String(), nullable=False),
        sa.Column("entity_id", sa.UUID(), nullable=False),
        sa.Column("tools", postgresql.JSONB(astext_type=sa.Text()), nullable=True),
        sa.Column("agent_preset_id", sa.UUID(), nullable=True),
        sa.Column("harness_type", sa.String(length=50), nullable=True),
        sa.Column("sdk_session_id", sa.String(length=120), nullable=True),
        sa.Column("curr_run_id", sa.UUID(), nullable=True),
        sa.Column("last_stream_id", sa.String(length=128), nullable=True),
        sa.Column("workspace_id", sa.UUID(), nullable=False),
        sa.Column("surrogate_id", sa.Integer(), nullable=False),
        sa.Column(
            "created_at",
            sa.TIMESTAMP(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.Column(
            "updated_at",
            sa.TIMESTAMP(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.ForeignKeyConstraint(
            ["agent_preset_id"],
            ["agent_preset.id"],
            name=op.f("fk_agent_session_agent_preset_id_agent_preset"),
            ondelete="SET NULL",
        ),
        sa.ForeignKeyConstraint(
            ["created_by"],
            ["user.id"],
            name=op.f("fk_agent_session_created_by_user"),
            ondelete="CASCADE",
        ),
        sa.ForeignKeyConstraint(
            ["workspace_id"],
            ["workspace.id"],
            name=op.f("fk_agent_session_workspace_id_workspace"),
            ondelete="CASCADE",
        ),
        sa.PrimaryKeyConstraint("surrogate_id", name=op.f("pk_agent_session")),
    )
    op.create_index(
        op.f("ix_agent_session_curr_run_id"),
        "agent_session",
        ["curr_run_id"],
        unique=False,
    )
    op.create_index(
        op.f("ix_agent_session_entity_type"),
        "agent_session",
        ["entity_type"],
        unique=False,
    )
    op.create_index(op.f("ix_agent_session_id"), "agent_session", ["id"], unique=True)
    op.create_table(
        "agent_session_history",
        sa.Column("id", sa.UUID(), nullable=False),
        sa.Column("session_id", sa.UUID(), nullable=False),
        sa.Column("content", postgresql.JSONB(astext_type=sa.Text()), nullable=False),
        sa.Column("kind", sa.String(length=50), nullable=False),
        sa.Column("workspace_id", sa.UUID(), nullable=False),
        sa.Column("surrogate_id", sa.Integer(), nullable=False),
        sa.Column(
            "created_at",
            sa.TIMESTAMP(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.Column(
            "updated_at",
            sa.TIMESTAMP(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.ForeignKeyConstraint(
            ["session_id"],
            ["agent_session.id"],
            name=op.f("fk_agent_session_history_session_id_agent_session"),
            ondelete="CASCADE",
        ),
        sa.ForeignKeyConstraint(
            ["workspace_id"],
            ["workspace.id"],
            name=op.f("fk_agent_session_history_workspace_id_workspace"),
            ondelete="CASCADE",
        ),
        sa.PrimaryKeyConstraint("surrogate_id", name=op.f("pk_agent_session_history")),
    )
    op.create_index(
        op.f("ix_agent_session_history_id"),
        "agent_session_history",
        ["id"],
        unique=True,
    )
    op.create_index(
        op.f("ix_agent_session_history_kind"),
        "agent_session_history",
        ["kind"],
        unique=False,
    )
    op.create_index(
        op.f("ix_agent_session_history_session_id"),
        "agent_session_history",
        ["session_id"],
        unique=False,
    )
    # Backfill placeholder sessions for pre-existing approval rows so the FK can
    # be added safely on databases that created approvals before agent_session.
    op.execute(
        sa.text(
            """
            INSERT INTO agent_session (
                id,
                title,
                created_by,
                entity_type,
                entity_id,
                tools,
                agent_preset_id,
                harness_type,
                curr_run_id,
                last_stream_id,
                workspace_id
            )
            SELECT
                missing.session_id,
                'Migrated approval session',
                NULL,
                'workflow',
                missing.session_id,
                NULL::jsonb,
                NULL,
                NULL,
                missing.session_id,
                NULL,
                missing.workspace_id
            FROM (
                SELECT DISTINCT a.session_id, a.workspace_id
                FROM approval AS a
                LEFT JOIN agent_session AS s
                    ON s.id = a.session_id
                WHERE a.session_id IS NOT NULL
                  AND s.id IS NULL
            ) AS missing
            """
        )
    )
    op.alter_column("approval", "session_id", existing_type=sa.UUID(), nullable=True)
    op.create_foreign_key(
        op.f("fk_approval_session_id_agent_session"),
        "approval",
        "agent_session",
        ["session_id"],
        ["id"],
        ondelete="SET NULL",
    )
    op.alter_column(
        "registry_version",
        "tarball_uri",
        existing_type=sa.VARCHAR(),
        comment=None,
        existing_comment="S3 URI to the compressed tarball venv (for git registries)",
        existing_nullable=False,
    )
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.alter_column(
        "registry_version",
        "tarball_uri",
        existing_type=sa.VARCHAR(),
        comment="S3 URI to the compressed tarball venv (for git registries)",
        existing_nullable=False,
    )
    op.drop_constraint(
        op.f("fk_approval_session_id_agent_session"), "approval", type_="foreignkey"
    )
    op.alter_column("approval", "session_id", existing_type=sa.UUID(), nullable=False)
    op.drop_index(
        op.f("ix_agent_session_history_session_id"), table_name="agent_session_history"
    )
    op.drop_index(
        op.f("ix_agent_session_history_kind"), table_name="agent_session_history"
    )
    op.drop_index(
        op.f("ix_agent_session_history_id"), table_name="agent_session_history"
    )
    op.drop_table("agent_session_history")
    op.drop_index(op.f("ix_agent_session_id"), table_name="agent_session")
    op.drop_index(op.f("ix_agent_session_entity_type"), table_name="agent_session")
    op.drop_index(op.f("ix_agent_session_curr_run_id"), table_name="agent_session")
    op.drop_table("agent_session")
    # ### end Alembic commands ###
