#!/usr/bin/env python3

from sqlalchemy import create_engine, String, update, bindparam, case
from sqlalchemy.orm import (
    DeclarativeBase,
    Mapped,
    mapped_column,
    sessionmaker,
    scoped_session,
)
from typing import Generator
from contextlib import contextmanager

db_url = "sqlite:////tmp/rhys.db"


class Base(DeclarativeBase):
    pass


class Rhys(Base):
    __tablename__ = "rhys"

    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
    ref_number: Mapped[int] = mapped_column()
    name: Mapped[str] = mapped_column(String(256), index=True)

    def __init__(self, session, ref_number: int, name: str) -> None:
        self.ref_number = ref_number
        self.name = name
        session.add(self)
        session.flush()


@contextmanager
def Session() -> Generator[scoped_session, None, None]:
    Session = scoped_session(sessionmaker(bind=engine))
    yield Session
    Session.commit()
    Session.close()


engine = create_engine(db_url)
Base.metadata.create_all(engine)

inital_number_of_records = 10


def move_rows(session):
    new_row = 6

    with Session() as session:
        # new_record = Rhys(session, new_row, f"new {new_row=}")
        # Move rows

        stmt = (
            update(Rhys)
            .where(Rhys.ref_number > new_row)
            # .where(Rhys.id.in_(session.query(Rhys.id).order_by(Rhys.id.desc())))
            .values({Rhys.ref_number: Rhys.ref_number + 1})
        )

        session.execute(stmt)


sqla_map = []
for k, v in zip(range(11), [0, 1, 2, 3, 4, 7, 8, 10, 5, 6, 9]):
    sqla_map.append({"oldrow": k, "newrow": v})

# for a, b in sqla_map.items():
#     print(f"{a} > {b}")

with Session() as session:
    for a in range(inital_number_of_records):
        _ = Rhys(session, a, f"record: {a}")

    stmt = update(Rhys).values(
        ref_number=case(
            {item['oldrow']: item['newrow'] for item in sqla_map},
            value=Rhys.ref_number
        )
    )

    session.connection().execute(stmt, sqla_map)