Model/backend/app/db/functions/magic_plan_functions.py

from typing import Any, cast

from sqlalchemy import delete, select
from sqlalchemy.dialects.postgresql import insert as pg_insert
from sqlmodel import Session, col

from datatypes.magicplan.domain.models import Floor, Plan
from backend.app.db.models.magic_plan import (
    MagicPlanDoorModel,
    MagicPlanFloorModel,
    MagicPlanPlanModel,
    MagicPlanRoomModel,
    MagicPlanWindowModel,
)


def save_plan(session: Session, plan: Plan) -> None:
    plan_id: int = _upsert_plan(session, plan)
    _delete_children(session, plan_id)
    floor_ids: list[int] = _insert_floors(session, plan.floors, plan_id)
    room_ids: list[int] = _insert_rooms(session, plan.floors, floor_ids)
    _insert_windows_and_doors(session, plan.floors, room_ids)


def _upsert_plan(session: Session, plan: Plan) -> int:
    stmt = (
        pg_insert(MagicPlanPlanModel)
        .values(
            magic_plan_uid=plan.uid,
            name=plan.name,
            address=plan.address,
            postcode=plan.postcode,
        )
        .on_conflict_do_update(
            index_elements=["magic_plan_uid"],
            set_={
                "name": plan.name,
                "address": plan.address,
                "postcode": plan.postcode,
            },
        )
        .returning(col(MagicPlanPlanModel.id))
    )
    row_id: int = session.execute(stmt).scalar_one()
    return row_id


def _delete_children(session: Session, plan_id: int) -> None:
    floor_subq = (
        select(col(MagicPlanFloorModel.id))
        .where(col(MagicPlanFloorModel.magic_plan_plan_id) == plan_id)
        .scalar_subquery()
    )
    room_subq = (
        select(col(MagicPlanRoomModel.id))
        .where(col(MagicPlanRoomModel.magic_plan_floor_id).in_(floor_subq))
        .scalar_subquery()
    )
    session.execute(
        delete(MagicPlanWindowModel).where(
            col(MagicPlanWindowModel.magic_plan_room_id).in_(room_subq)
        )
    )
    session.execute(
        delete(MagicPlanDoorModel).where(
            col(MagicPlanDoorModel.magic_plan_room_id).in_(room_subq)
        )
    )
    session.execute(
        delete(MagicPlanRoomModel).where(
            col(MagicPlanRoomModel.magic_plan_floor_id).in_(floor_subq)
        )
    )
    session.execute(
        delete(MagicPlanFloorModel).where(
            col(MagicPlanFloorModel.magic_plan_plan_id) == plan_id
        )
    )


def _insert_floors(session: Session, floors: list[Floor], plan_id: int) -> list[int]:
    rows: list[dict[str, Any]] = [
        {"magic_plan_plan_id": plan_id, "level": floor.level} for floor in floors
    ]
    result = session.execute(
        pg_insert(MagicPlanFloorModel)
        .values(rows)
        .returning(col(MagicPlanFloorModel.id))
    )
    return cast(list[int], list(result.scalars().all()))


def _insert_rooms(
    session: Session, floors: list[Floor], floor_ids: list[int]
) -> list[int]:
    rows: list[dict[str, Any]] = [
        {
            "magic_plan_floor_id": floor_id,
            "name": room.name,
            "width_m": room.width_m,
            "length_m": room.length_m,
            "area_m2": room.area_m2,
        }
        for floor, floor_id in zip(floors, floor_ids)
        for room in floor.rooms
    ]
    result = session.execute(
        pg_insert(MagicPlanRoomModel).values(rows).returning(col(MagicPlanRoomModel.id))
    )
    return cast(list[int], list(result.scalars().all()))


def _insert_windows_and_doors(
    session: Session, floors: list[Floor], room_ids: list[int]
) -> None:
    all_rooms = [room for floor in floors for room in floor.rooms]

    window_rows: list[dict[str, Any]] = [
        {
            "magic_plan_room_id": room_id,
            "width_m": window.width_m,
            "height_m": window.height_m,
            "area_m2": window.area_m2,
            "opening_type": window.opening_type,
        }
        for room, room_id in zip(all_rooms, room_ids)
        for window in room.windows
    ]
    door_rows: list[dict[str, Any]] = [
        {
            "magic_plan_room_id": room_id,
            "width_mm": door.width_mm,
        }
        for room, room_id in zip(all_rooms, room_ids)
        for door in room.doors
    ]

    if window_rows:
        session.execute(pg_insert(MagicPlanWindowModel).values(window_rows))
    if door_rows:
        session.execute(pg_insert(MagicPlanDoorModel).values(door_rows))