Model/repositories/magic_plan/magic_plan_postgres_repository.py

from __future__ import annotations

from typing import Any, Optional, 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 (
    Door,
    DoorVentilation,
    Floor,
    Plan,
    Room,
    Window,
    WindowVentilation,
)
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 get_plan_by_uploaded_file_id(self, uploaded_file_id: int) -> Optional[Plan]:
        plan_row = self._session.execute(  # pyright: ignore[reportDeprecated]
            select(MagicPlanPlanModel).where(
                col(MagicPlanPlanModel.uploaded_file_id) == uploaded_file_id
            )
        ).scalars().one_or_none()
        if plan_row is None:
            return None

        plan_id = cast(int, plan_row.id)

        floor_rows = list(
            self._session.execute(  # pyright: ignore[reportDeprecated]
                select(MagicPlanFloorModel).where(
                    col(MagicPlanFloorModel.magic_plan_plan_id) == plan_id
                )
            ).scalars().all()
        )
        floor_ids = [cast(int, f.id) for f in floor_rows]

        room_rows = list(
            self._session.execute(  # pyright: ignore[reportDeprecated]
                select(MagicPlanRoomModel).where(
                    col(MagicPlanRoomModel.magic_plan_floor_id).in_(floor_ids)
                )
            ).scalars().all()
        )
        room_ids = [cast(int, r.id) for r in room_rows]

        window_rows = list(
            self._session.execute(  # pyright: ignore[reportDeprecated]
                select(MagicPlanWindowModel).where(
                    col(MagicPlanWindowModel.magic_plan_room_id).in_(room_ids)
                )
            ).scalars().all()
        )
        window_ids = [cast(int, w.id) for w in window_rows]

        door_rows = list(
            self._session.execute(  # pyright: ignore[reportDeprecated]
                select(MagicPlanDoorModel).where(
                    col(MagicPlanDoorModel.magic_plan_room_id).in_(room_ids)
                )
            ).scalars().all()
        )
        door_ids = [cast(int, d.id) for d in door_rows]

        window_vent_by_window_id = {
            cast(int, wv.magic_plan_window_id): wv
            for wv in self._session.execute(  # pyright: ignore[reportDeprecated]
                select(MagicPlanWindowVentilationModel).where(
                    col(MagicPlanWindowVentilationModel.magic_plan_window_id).in_(
                        window_ids
                    )
                )
            ).scalars().all()
        }
        door_vent_by_door_id = {
            cast(int, dv.magic_plan_door_id): dv
            for dv in self._session.execute(  # pyright: ignore[reportDeprecated]
                select(MagicPlanDoorVentilationModel).where(
                    col(MagicPlanDoorVentilationModel.magic_plan_door_id).in_(door_ids)
                )
            ).scalars().all()
        }

        windows_by_room_id: dict[int, list[Window]] = {}
        for w_row, w_id in zip(window_rows, window_ids):
            room_id = cast(int, w_row.magic_plan_room_id)
            wv_row = window_vent_by_window_id.get(w_id)
            ventilation = (
                WindowVentilation(
                    opening_type=wv_row.opening_type,
                    num_openings=wv_row.num_openings,
                    pct_openable=wv_row.pct_openable,
                    trickle_vent_area_mm2=wv_row.trickle_vent_area_mm2,
                    num_trickle_vents=wv_row.num_trickle_vents,
                )
                if wv_row is not None
                else None
            )
            windows_by_room_id.setdefault(room_id, []).append(
                Window(
                    width_m=cast(float, w_row.width_m),
                    height_m=cast(float, w_row.height_m),
                    area_m2=cast(float, w_row.area_m2),
                    ventilation=ventilation,
                )
            )

        doors_by_room_id: dict[int, list[Door]] = {}
        for d_row, d_id in zip(door_rows, door_ids):
            room_id = cast(int, d_row.magic_plan_room_id)
            dv_row = door_vent_by_door_id.get(d_id)
            ventilation = (
                DoorVentilation(undercut_mm=dv_row.undercut_mm)
                if dv_row is not None
                else None
            )
            doors_by_room_id.setdefault(room_id, []).append(
                Door(
                    width_mm=cast(float, d_row.width_mm),
                    height_mm=cast(float, d_row.height_mm),
                    ventilation=ventilation,
                )
            )

        rooms_by_floor_id: dict[int, list[Room]] = {}
        for r_row, r_id in zip(room_rows, room_ids):
            floor_id = cast(int, r_row.magic_plan_floor_id)
            rooms_by_floor_id.setdefault(floor_id, []).append(
                Room(
                    name=cast(str, r_row.name),
                    width_m=cast(float, r_row.width_m),
                    length_m=cast(float, r_row.length_m),
                    area_m2=cast(float, r_row.area_m2),
                    windows=windows_by_room_id.get(r_id, []),
                    doors=doors_by_room_id.get(r_id, []),
                )
            )

        floors = [
            Floor(
                level=f_row.level,
                name=None,
                rooms=rooms_by_floor_id.get(cast(int, f_row.id), []),
            )
            for f_row in floor_rows
        ]

        return Plan(
            uid=cast(str, plan_row.magic_plan_uid),
            name=plan_row.name,
            address=plan_row.address,
            postcode=plan_row.postcode,
            floors=floors,
        )

    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)
            )