"""add_agent_ids_col_to_session

Revision ID: 69c059996cbd
Revises: 9e599b62f6f1
Create Date: 2023-07-05 23:39:21.631462

"""

import uuid
from collections import defaultdict
from typing import Any

import sqlalchemy as sa
from alembic import op
from sqlalchemy.orm import registry
from sqlalchemy.sql.expression import bindparam

from ai.backend.manager.models.base import GUID, convention

# revision identifiers, used by Alembic.
revision = "69c059996cbd"
down_revision = "9e599b62f6f1"
branch_labels = None
depends_on = None

metadata = sa.MetaData(naming_convention=convention)
mapper_registry = registry(metadata=metadata)
Base = mapper_registry.generate_base()

PAGE_SIZE = 100


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    conn = op.get_bind()
    op.add_column("sessions", sa.Column("agent_ids", sa.ARRAY(sa.String()), nullable=True))

    kernels = sa.Table(
        "kernels",
        metadata,
        sa.Column("id", GUID, primary_key=True),
        sa.Column("session_id", GUID, nullable=False),
        sa.Column("agent", sa.String(length=64), sa.ForeignKey("agents.id"), nullable=True),
    )

    class SessionRow(Base):  # type: ignore[valid-type, misc]
        __tablename__ = "sessions"
        id: sa.Column[GUID[uuid.UUID]] = sa.Column(
            "id",
            GUID(),
            server_default=sa.text("uuid_generate_v4()"),
            nullable=False,
            primary_key=True,
        )
        agent_ids: sa.Column[list[str] | None] = sa.Column(
            "agent_ids", sa.ARRAY(sa.String), nullable=True
        )

    kernel_cnt: int = conn.execute(sa.select(sa.func.count()).select_from(kernels)).scalar() or 0

    for offset in range(0, kernel_cnt, PAGE_SIZE):
        session_agent_ids_map: dict[Any, set[Any]] = defaultdict(set)
        session_id_query = (
            sa.select(sa.distinct(kernels.c.session_id))
            .where(kernels.c.agent.is_not(None))
            .order_by(kernels.c.session_id)
            .offset(offset)
            .limit(PAGE_SIZE)
        )
        session_ids = conn.execute(session_id_query).scalars().all()
        stmt = sa.select(kernels).where(kernels.c.session_id.in_(session_ids))
        kernel_rows = conn.execute(stmt).fetchall()
        for row in kernel_rows:
            session_agent_ids_map[row.session_id].add(row.agent)
        if session_agent_ids_map:
            conn.execute(
                sa.update(SessionRow)
                .values({"agent_ids": bindparam("agent_ids")})
                .where(SessionRow.id == bindparam("b_session_id")),
                [
                    {
                        "b_session_id": session_id,
                        "agent_ids": list(agent_ids),
                    }
                    for session_id, agent_ids in session_agent_ids_map.items()
                ],
            )
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_column("sessions", "agent_ids")
    # ### end Alembic commands ###
