Model/repositories/magic_plan/magic_plan_postgres_repository.py

from __future__ import annotations

from typing import Any, NamedTuple, 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 _Rows(NamedTuple):
    floors: list[MagicPlanFloorModel]
    rooms: list[MagicPlanRoomModel]
    windows: list[MagicPlanWindowModel]
    doors: list[MagicPlanDoorModel]
    win_vents: list[MagicPlanWindowVentilationModel]
    door_vents: list[MagicPlanDoorVentilationModel]


def _build_windows(
    rows: list[MagicPlanWindowModel],
    vents: list[MagicPlanWindowVentilationModel],
) -> dict[int, list[Window]]:
    vent_by_id = {wv.magic_plan_window_id: wv for wv in vents}
    result: dict[int, list[Window]] = {}
    for row in rows:
        wv = vent_by_id.get(cast(int, row.id))
        result.setdefault(row.magic_plan_room_id, []).append(
            Window(
                width_m=cast(float, row.width_m),
                height_m=cast(float, row.height_m),
                area_m2=cast(float, row.area_m2),
                ventilation=WindowVentilation(
                    opening_type=wv.opening_type,
                    num_openings=wv.num_openings,
                    pct_openable=wv.pct_openable,
                    trickle_vent_area_mm2=wv.trickle_vent_area_mm2,
                    num_trickle_vents=wv.num_trickle_vents,
                ) if wv else None,
            )
        )
    return result


def _build_doors(
    rows: list[MagicPlanDoorModel],
    vents: list[MagicPlanDoorVentilationModel],
) -> dict[int, list[Door]]:
    vent_by_id = {dv.magic_plan_door_id: dv for dv in vents}
    result: dict[int, list[Door]] = {}
    for row in rows:
        dv = vent_by_id.get(cast(int, row.id))
        result.setdefault(row.magic_plan_room_id, []).append(
            Door(
                width_mm=cast(float, row.width_mm),
                height_mm=cast(float, row.height_mm),
                ventilation=DoorVentilation(undercut_mm=dv.undercut_mm) if dv else None,
            )
        )
    return result


def _build_rooms(
    rows: list[MagicPlanRoomModel],
    windows_by_room: dict[int, list[Window]],
    doors_by_room: dict[int, list[Door]],
) -> dict[int, list[Room]]:
    result: dict[int, list[Room]] = {}
    for row in rows:
        room_id = cast(int, row.id)
        result.setdefault(row.magic_plan_floor_id, []).append(
            Room(
                name=cast(str, row.name),
                width_m=cast(float, row.width_m),
                length_m=cast(float, row.length_m),
                area_m2=cast(float, row.area_m2),
                windows=windows_by_room.get(room_id, []),
                doors=doors_by_room.get(room_id, []),
            )
        )
    return result


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._fetch_one(
            select(MagicPlanPlanModel).where(
                col(MagicPlanPlanModel.uploaded_file_id) == uploaded_file_id
            )
        )
        if plan_row is None:
            return None

        rows = self._fetch_rows(cast(int, plan_row.id))
        windows_by_room = _build_windows(rows.windows, rows.win_vents)
        doors_by_room = _build_doors(rows.doors, rows.door_vents)
        rooms_by_floor = _build_rooms(rows.rooms, windows_by_room, doors_by_room)

        return Plan(
            uid=cast(str, plan_row.magic_plan_uid),
            name=plan_row.name,
            address=plan_row.address,
            postcode=plan_row.postcode,
            floors=[
                Floor(level=f.level, name=None, rooms=rooms_by_floor.get(cast(int, f.id), []))
                for f in rows.floors
            ],
        )

    def _fetch_rows(self, plan_id: int) -> _Rows:
        floor_rows: list[MagicPlanFloorModel] = self._fetch_many(
            select(MagicPlanFloorModel).where(
                col(MagicPlanFloorModel.magic_plan_plan_id) == plan_id
            )
        )
        floor_ids = [cast(int, f.id) for f in floor_rows]
        room_rows: list[MagicPlanRoomModel] = self._fetch_many(
            select(MagicPlanRoomModel).where(
                col(MagicPlanRoomModel.magic_plan_floor_id).in_(floor_ids)
            )
        )
        room_ids = [cast(int, r.id) for r in room_rows]
        window_rows: list[MagicPlanWindowModel] = self._fetch_many(
            select(MagicPlanWindowModel).where(
                col(MagicPlanWindowModel.magic_plan_room_id).in_(room_ids)
            )
        )
        door_rows: list[MagicPlanDoorModel] = self._fetch_many(
            select(MagicPlanDoorModel).where(
                col(MagicPlanDoorModel.magic_plan_room_id).in_(room_ids)
            )
        )
        win_vents: list[MagicPlanWindowVentilationModel] = self._fetch_many(
            select(MagicPlanWindowVentilationModel).where(
                col(MagicPlanWindowVentilationModel.magic_plan_window_id).in_(
                    [cast(int, w.id) for w in window_rows]
                )
            )
        )
        door_vents: list[MagicPlanDoorVentilationModel] = self._fetch_many(
            select(MagicPlanDoorVentilationModel).where(
                col(MagicPlanDoorVentilationModel.magic_plan_door_id).in_(
                    [cast(int, d.id) for d in door_rows]
                )
            )
        )
        return _Rows(floor_rows, room_rows, window_rows, door_rows, win_vents, door_vents)

    def _fetch_one(self, stmt: Any) -> Any:
        return self._session.execute(stmt).scalars().one_or_none()  # pyright: ignore[reportDeprecated]

    def _fetch_many(self, stmt: Any) -> Any:
        return list(self._session.execute(stmt).scalars().all())  # pyright: ignore[reportDeprecated]

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