Model/repositories/magic_plan/magic_plan_postgres_repository.py

from __future__ import annotations

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 domain.magicplan.models import Floor, Plan
from infrastructure.postgres.magic_plan_tables import (
    MagicPlanDoorModel,
    MagicPlanDoorVentilationModel,
    MagicPlanFloorModel,
    MagicPlanPlanModel,
    MagicPlanRoomModel,
    MagicPlanWindowModel,
    MagicPlanWindowVentilationModel,
)
from repositories.magic_plan.magic_plan_repository import MagicPlanRepository


class MagicPlanPostgresRepository(MagicPlanRepository):
    def __init__(self, session: Session) -> None:
        self._session = session

    def save(self, plan: Plan, uploaded_file_id: int) -> None:
        plan_id = self._upsert_plan(plan, uploaded_file_id)
        self._delete_children(plan_id)
        floor_ids = self._insert_floors(plan.floors, plan_id)
        room_ids = self._insert_rooms(plan.floors, floor_ids)
        window_ids, door_ids = self._insert_windows_and_doors(plan.floors, room_ids)
        self._insert_ventilation(plan.floors, window_ids, door_ids)

    def _upsert_plan(self, plan: Plan, uploaded_file_id: int) -> int:
        row_data: dict[str, Any] = MagicPlanPlanModel.from_domain(
            plan, uploaded_file_id
        ).model_dump(exclude={"id"})
        stmt = (
            pg_insert(MagicPlanPlanModel)
            .values(**row_data)
            .on_conflict_do_update(
                index_elements=["magic_plan_uid"],
                set_={k: v for k, v in row_data.items() if k != "magic_plan_uid"},
            )
            .returning(col(MagicPlanPlanModel.id))
        )
        return cast(
            int, self._session.execute(stmt).scalar_one()
        )  # pyright: ignore[reportDeprecated]

    def _delete_children(self, 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()
        )
        window_subq = (
            select(col(MagicPlanWindowModel.id))
            .where(col(MagicPlanWindowModel.magic_plan_room_id).in_(room_subq))
            .scalar_subquery()
        )
        door_subq = (
            select(col(MagicPlanDoorModel.id))
            .where(col(MagicPlanDoorModel.magic_plan_room_id).in_(room_subq))
            .scalar_subquery()
        )
        self._session.execute(  # pyright: ignore[reportDeprecated]
            delete(MagicPlanWindowVentilationModel).where(
                col(MagicPlanWindowVentilationModel.magic_plan_window_id).in_(
                    window_subq
                )
            )
        )
        self._session.execute(  # pyright: ignore[reportDeprecated]
            delete(MagicPlanDoorVentilationModel).where(
                col(MagicPlanDoorVentilationModel.magic_plan_door_id).in_(door_subq)
            )
        )
        self._session.execute(  # pyright: ignore[reportDeprecated]
            delete(MagicPlanWindowModel).where(
                col(MagicPlanWindowModel.magic_plan_room_id).in_(room_subq)
            )
        )
        self._session.execute(  # pyright: ignore[reportDeprecated]
            delete(MagicPlanDoorModel).where(
                col(MagicPlanDoorModel.magic_plan_room_id).in_(room_subq)
            )
        )
        self._session.execute(  # pyright: ignore[reportDeprecated]
            delete(MagicPlanRoomModel).where(
                col(MagicPlanRoomModel.magic_plan_floor_id).in_(floor_subq)
            )
        )
        self._session.execute(  # pyright: ignore[reportDeprecated]
            delete(MagicPlanFloorModel).where(
                col(MagicPlanFloorModel.magic_plan_plan_id) == plan_id
            )
        )

    def _insert_floors(self, floors: list[Floor], plan_id: int) -> list[int]:
        rows: list[dict[str, Any]] = [
            MagicPlanFloorModel.from_domain(floor, plan_id).model_dump(exclude={"id"})
            for floor in floors
        ]
        result = self._session.execute(  # pyright: ignore[reportDeprecated]
            pg_insert(MagicPlanFloorModel)
            .values(rows)
            .returning(col(MagicPlanFloorModel.id))
        )
        return cast(list[int], list(result.scalars().all()))

    def _insert_rooms(self, floors: list[Floor], floor_ids: list[int]) -> list[int]:
        rows: list[dict[str, Any]] = [
            MagicPlanRoomModel.from_domain(room, floor_id).model_dump(exclude={"id"})
            for floor, floor_id in zip(floors, floor_ids)
            for room in floor.rooms
        ]
        result = self._session.execute(  # pyright: ignore[reportDeprecated]
            pg_insert(MagicPlanRoomModel)
            .values(rows)
            .returning(col(MagicPlanRoomModel.id))
        )
        return cast(list[int], list(result.scalars().all()))

    def _insert_windows_and_doors(
        self, floors: list[Floor], room_ids: list[int]
    ) -> tuple[list[int], list[int]]:
        all_rooms = [room for floor in floors for room in floor.rooms]
        window_rows: list[dict[str, Any]] = [
            MagicPlanWindowModel.from_domain(window, room_id).model_dump(exclude={"id"})
            for room, room_id in zip(all_rooms, room_ids)
            for window in room.windows
        ]
        door_rows: list[dict[str, Any]] = [
            MagicPlanDoorModel.from_domain(door, room_id).model_dump(exclude={"id"})
            for room, room_id in zip(all_rooms, room_ids)
            for door in room.doors
        ]
        window_ids: list[int] = []
        door_ids: list[int] = []
        if window_rows:
            result = self._session.execute(  # pyright: ignore[reportDeprecated]
                pg_insert(MagicPlanWindowModel)
                .values(window_rows)
                .returning(col(MagicPlanWindowModel.id))
            )
            window_ids = cast(list[int], list(result.scalars().all()))
        if door_rows:
            result = self._session.execute(  # pyright: ignore[reportDeprecated]
                pg_insert(MagicPlanDoorModel)
                .values(door_rows)
                .returning(col(MagicPlanDoorModel.id))
            )
            door_ids = cast(list[int], list(result.scalars().all()))
        return window_ids, door_ids

    def _insert_ventilation(
        self,
        floors: list[Floor],
        window_ids: list[int],
        door_ids: list[int],
    ) -> None:
        all_rooms = [room for floor in floors for room in floor.rooms]
        all_windows = [w for room in all_rooms for w in room.windows]
        all_doors = [d for room in all_rooms for d in room.doors]

        window_vent_rows: list[dict[str, Any]] = [
            MagicPlanWindowVentilationModel.from_domain(w.ventilation, wid).model_dump(
                exclude={"id"}
            )
            for w, wid in zip(all_windows, window_ids)
            if w.ventilation is not None
        ]
        door_vent_rows: list[dict[str, Any]] = [
            MagicPlanDoorVentilationModel.from_domain(d.ventilation, did).model_dump(
                exclude={"id"}
            )
            for d, did in zip(all_doors, door_ids)
            if d.ventilation is not None
        ]
        if window_vent_rows:
            self._session.execute(  # pyright: ignore[reportDeprecated]
                pg_insert(MagicPlanWindowVentilationModel).values(window_vent_rows)
            )
        if door_vent_rows:
            self._session.execute(  # pyright: ignore[reportDeprecated]
                pg_insert(MagicPlanDoorVentilationModel).values(door_vent_rows)
            )