"""add_images_col_to_session

Revision ID: eb9441fcf90a
Revises: 69c059996cbd
Create Date: 2023-07-06 13:51:52.098587

"""

from __future__ import annotations

from collections import defaultdict
from typing import Any, cast
from uuid import UUID

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 = "eb9441fcf90a"
down_revision = "69c059996cbd"
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
DEFAULT_ROLE = "main"


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    conn = op.get_bind()
    op.add_column("sessions", sa.Column("images", 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("image", sa.String(length=512)),
        sa.Column(
            "cluster_role",
            sa.String(length=16),
        ),
    )

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

    kernel_cnt_result = conn.execute(sa.select(sa.func.count()).select_from(kernels)).scalar()
    kernel_cnt = kernel_cnt_result if kernel_cnt_result is not None else 0

    for offset in range(0, kernel_cnt, PAGE_SIZE):
        session_agent_ids_map: defaultdict[UUID, list[Any]] = defaultdict(list)
        session_id_query = (
            sa.select(sa.distinct(kernels.c.session_id))
            .where(kernels.c.image.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:
            row_dict = cast(dict[str, Any], row._mapping)
            img_list = session_agent_ids_map[row_dict["session_id"]]
            if row_dict["image"] in img_list:
                continue
            if row_dict["cluster_role"] == DEFAULT_ROLE:
                img_list.insert(0, row_dict["image"])
            else:
                img_list.append(row_dict["image"])
        if session_agent_ids_map:
            conn.execute(
                sa.update(SessionRow)
                .values({"images": bindparam("images")})
                .where(SessionRow.id == bindparam("b_session_id")),
                [
                    {
                        "b_session_id": session_id,
                        "images": list(images),
                    }
                    for session_id, images in session_agent_ids_map.items()
                ],
            )
    # ### end Alembic commands ###


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