Model/repositories/product/product_postgres_repository.py

from __future__ import annotations

from typing import Optional

from sqlmodel import Session, col, select

from domain.modelling.contingencies import contingency_rate
from domain.modelling.product import Product
from infrastructure.postgres.product_table import MaterialRow
from repositories.product.product_repository import ProductNotFound, ProductRepository


# The domain ``MeasureType`` vocabulary and the catalogue's ``material.type``
# pgEnum drifted apart: these five measures are spelled differently on the
# catalogue side (and querying the domain spelling raises a pgEnum DataError
# that poisons the session's transaction). Translate them to the catalogue's
# own vocabulary at this boundary so the domain enum stays stable. Every other
# MeasureType already matches its material.type and maps to itself.
_MATERIAL_TYPE_BY_MEASURE: dict[str, str] = {
    "low_energy_lighting": "low_energy_lighting_installation",
    "gas_boiler_upgrade": "boiler_upgrade",
    "system_tune_up": "roomstat_programmer_trvs",
    "system_tune_up_zoned": "time_temperature_zone_control",
    "sloping_ceiling_insulation": "room_roof_insulation",
}


class ProductPostgresRepository(ProductRepository):
    """Reads the ``material`` catalogue table and maps an active row to a
    Product: `total_cost` becomes the fully-loaded `unit_cost_per_m2`, and the
    per-Measure-Type contingency is joined from config."""

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

    def get(self, measure_type: str) -> Product:
        # Resolve the domain MeasureType to the catalogue's ``material.type``
        # spelling (identity for all but the five drifted types above).
        catalogue_type = _MATERIAL_TYPE_BY_MEASURE.get(measure_type, measure_type)
        # The live catalogue holds many active rows per type; order by id so the
        # pick is deterministic (a re-seed prices the same) rather than relying
        # on the database's physical row order.
        row: MaterialRow | None = self._session.exec(
            select(MaterialRow)
            .where(
                col(MaterialRow.type) == catalogue_type,
                col(MaterialRow.is_active).is_(True),
            )
            .order_by(col(MaterialRow.id))
        ).first()
        if row is None:
            raise ProductNotFound(
                f"no active product for measure type {measure_type!r}"
            )
        if row.total_cost is None:
            raise ValueError(f"product {measure_type!r} has no total_cost")
        return Product(
            measure_type=measure_type,
            unit_cost_per_m2=row.total_cost,
            contingency_rate=contingency_rate(measure_type),
            id=row.id,
        )


class MaterialSnapshotRepository(ProductRepository):
    """An in-memory snapshot of the active ``material`` catalogue, read once from
    a Session via ``load``.

    It lets the Modelling stage price Measure Options *after* the read Session is
    closed, so a Lambda invocation can hold a single DB connection at a time — the
    read Session, then each per-Property write Unit of Work, never overlapping —
    instead of keeping the read Session's connection checked out across the whole
    write loop while the catalogue is queried lazily.

    ``get`` mirrors ``ProductPostgresRepository.get`` exactly: the same
    domain→``material.type`` drift mapping, the lowest-id active row per type, and
    the same ``ProductNotFound`` / missing-``total_cost`` errors. Because the
    snapshot is a plain dict lookup (no live query), the off-catalogue measures
    that would poison a live session's transaction simply miss and raise the
    benign ``ProductNotFound`` — the ``CompositeProductRepository`` override still
    keeps them off the catalogue first.
    """

    def __init__(
        self, rows_by_catalogue_type: dict[str, tuple[Optional[float], int]]
    ) -> None:
        self._rows_by_catalogue_type = rows_by_catalogue_type

    @classmethod
    def load(cls, session: Session) -> "MaterialSnapshotRepository":
        rows = session.exec(
            select(MaterialRow)
            .where(col(MaterialRow.is_active).is_(True))
            .order_by(col(MaterialRow.id))
        ).all()
        rows_by_catalogue_type: dict[str, tuple[Optional[float], int]] = {}
        for row in rows:
            # Lowest id wins per type — mirrors ``.first()`` after ``order_by(id)``.
            if row.type not in rows_by_catalogue_type:
                rows_by_catalogue_type[row.type] = (row.total_cost, row.id)
        return cls(rows_by_catalogue_type)

    def get(self, measure_type: str) -> Product:
        catalogue_type = _MATERIAL_TYPE_BY_MEASURE.get(measure_type, measure_type)
        entry = self._rows_by_catalogue_type.get(catalogue_type)
        if entry is None:
            raise ProductNotFound(
                f"no active product for measure type {measure_type!r}"
            )
        total_cost, row_id = entry
        if total_cost is None:
            raise ValueError(f"product {measure_type!r} has no total_cost")
        return Product(
            measure_type=measure_type,
            unit_cost_per_m2=total_cost,
            contingency_rate=contingency_rate(measure_type),
            id=row_id,
        )