Model/infrastructure/postgres/property_baseline_performance_table.py

from __future__ import annotations

from typing import ClassVar, Optional, cast, get_args

from sqlalchemy import Column
from sqlalchemy import Enum as SAEnum
from sqlmodel import Field, SQLModel

from datatypes.epc.domain.epc import Epc
from domain.billing.bill import Bill, BillSection, BillSectionCost
from domain.property_baseline.property_baseline_performance import PropertyBaselinePerformance
from domain.property_baseline.performance import Performance
from domain.property_baseline.rebaseliner import RebaselineReason

# Native-enum labels for the ``epc`` / ``rebaseline_reason`` columns, sourced
# from the domain types so the mirror can't drift from them.
_EPC_BANDS: tuple[str, ...] = tuple(band.value for band in Epc)
_REBASELINE_REASONS: tuple[str, ...] = get_args(RebaselineReason)

# Each Bill section's flat-column stem (``bill_{stem}_kwh`` / ``bill_{stem}_cost_gbp``).
_SECTION_COLUMN_STEM: dict[BillSection, str] = {
    BillSection.HEATING: "heating",
    BillSection.HOT_WATER: "hot_water",
    BillSection.LIGHTING: "lighting",
    BillSection.APPLIANCES: "appliances",
    BillSection.COOKING: "cooking",
    BillSection.PUMPS_FANS: "pumps_fans",
    BillSection.COOLING: "cooling",
}


class PropertyBaselinePerformanceModel(SQLModel, table=True):
    """The ``property_baseline_performance`` row — one per Property (ADR-0004).

    Flat typed columns (not a JSONB blob) so the FE can both surface the block
    and query the lodged-vs-effective pair. The production migration is FE-owned
    (Drizzle); see docs/migrations/property-baseline-performance-table.md.
    """

    __tablename__: ClassVar[str] = "property_baseline_performance"  # pyright: ignore[reportIncompatibleVariableOverride]

    id: Optional[int] = Field(default=None, primary_key=True)
    property_id: int = Field(unique=True, index=True)

    # ``*_epc_band`` and ``rebaseline_reason`` are native Postgres enums on the
    # live table (``epc`` / ``rebaseline_reason``, owned by Drizzle); declaring
    # them as bare ``str`` makes SQLAlchemy bind VARCHAR on INSERT, which
    # Postgres won't coerce to the enum. Bind the native enums explicitly
    # (``create_type=False``: the types already exist). Values stay plain
    # strings on the domain side (``Epc(...)`` / the RebaselineReason Literal).
    # The four lodged_* columns are nullable as a unit: a predicted Property has no
    # lodged cert, so it has no Lodged Performance (ADR-0004 amendment, #1361). They
    # are written all-set or all-NULL; lodged_sap_score is the read discriminator.
    lodged_sap_score: Optional[int] = None
    lodged_epc_band: Optional[str] = Field(
        default=None,
        sa_column=Column(SAEnum(*_EPC_BANDS, name="epc", create_type=False), nullable=True),
    )
    lodged_co2_emissions_t_per_yr: Optional[float] = None
    lodged_primary_energy_intensity_kwh_per_m2_yr: Optional[int] = None

    effective_sap_score: int
    effective_epc_band: str = Field(
        sa_column=Column(SAEnum(*_EPC_BANDS, name="epc", create_type=False), nullable=False)
    )
    effective_co2_emissions_t_per_yr: float
    effective_primary_energy_intensity_kwh_per_m2_yr: int

    rebaseline_reason: str = Field(
        sa_column=Column(
            SAEnum(*_REBASELINE_REASONS, name="rebaseline_reason", create_type=False),
            nullable=False,
        )
    )

    space_heating_kwh: float
    water_heating_kwh: float

    # The Fuel Rates snapshot period the bill was priced against (FE-owned column,
    # nullable). Not yet threaded through Bill Derivation, so left None for now.
    fuel_rates_period: Optional[str] = Field(default=None)

    # Bill Derivation block (ADR-0014 §6). Nullable: all None when no calculator
    # ran (stub path). Column names are unprefixed to mirror the FE-owned table —
    # the per-section ``heating_kwh`` / ``hot_water_kwh`` do not clash with the
    # recorded-demand ``space_heating_kwh`` / ``water_heating_kwh`` above.
    heating_kwh: Optional[float] = Field(default=None)
    heating_cost_gbp: Optional[float] = Field(default=None)
    hot_water_kwh: Optional[float] = Field(default=None)
    hot_water_cost_gbp: Optional[float] = Field(default=None)
    lighting_kwh: Optional[float] = Field(default=None)
    lighting_cost_gbp: Optional[float] = Field(default=None)
    appliances_kwh: Optional[float] = Field(default=None)
    appliances_cost_gbp: Optional[float] = Field(default=None)
    cooking_kwh: Optional[float] = Field(default=None)
    cooking_cost_gbp: Optional[float] = Field(default=None)
    pumps_fans_kwh: Optional[float] = Field(default=None)
    pumps_fans_cost_gbp: Optional[float] = Field(default=None)
    cooling_kwh: Optional[float] = Field(default=None)
    cooling_cost_gbp: Optional[float] = Field(default=None)
    standing_charges_gbp: Optional[float] = Field(default=None)
    seg_credit_gbp: Optional[float] = Field(default=None)
    total_annual_bill_gbp: Optional[float] = Field(default=None)

    @classmethod
    def from_domain(
        cls, baseline: PropertyBaselinePerformance, property_id: int
    ) -> "PropertyBaselinePerformanceModel":
        # A predicted Property has no Lodged Performance — the four lodged_* columns
        # are all NULL then (ADR-0004 amendment, #1361).
        lodged = baseline.lodged
        model = cls(
            property_id=property_id,
            lodged_sap_score=lodged.sap_score if lodged is not None else None,
            lodged_epc_band=lodged.epc_band.value if lodged is not None else None,
            lodged_co2_emissions_t_per_yr=lodged.co2_emissions if lodged is not None else None,
            lodged_primary_energy_intensity_kwh_per_m2_yr=(
                lodged.primary_energy_intensity if lodged is not None else None
            ),
            effective_sap_score=baseline.effective.sap_score,
            effective_epc_band=baseline.effective.epc_band.value,
            effective_co2_emissions_t_per_yr=baseline.effective.co2_emissions,
            effective_primary_energy_intensity_kwh_per_m2_yr=baseline.effective.primary_energy_intensity,
            rebaseline_reason=baseline.rebaseline_reason,
            space_heating_kwh=baseline.space_heating_kwh,
            water_heating_kwh=baseline.water_heating_kwh,
        )
        model._write_bill(baseline.bill)
        return model

    def _write_bill(self, bill: Optional[Bill]) -> None:
        """Flatten the Bill onto the ``bill_*`` columns. When ``bill`` is None
        (no calculator ran) every bill column is left None; a section absent from
        the mapping leaves its two columns None (None != 0 — it was not billed)."""
        if bill is None:
            return
        for section, stem in _SECTION_COLUMN_STEM.items():
            cost = bill.sections.get(section)
            setattr(self, f"{stem}_kwh", cost.kwh if cost is not None else None)
            setattr(
                self,
                f"{stem}_cost_gbp",
                cost.cost_gbp if cost is not None else None,
            )
        self.standing_charges_gbp = bill.standing_charges_gbp
        self.seg_credit_gbp = bill.seg_credit_gbp
        self.total_annual_bill_gbp = bill.total_gbp

    def to_domain(self) -> PropertyBaselinePerformance:
        return PropertyBaselinePerformance(
            lodged=self._read_lodged(),
            effective=Performance(
                sap_score=self.effective_sap_score,
                epc_band=Epc(self.effective_epc_band),
                co2_emissions=self.effective_co2_emissions_t_per_yr,
                primary_energy_intensity=self.effective_primary_energy_intensity_kwh_per_m2_yr,
            ),
            rebaseline_reason=cast(RebaselineReason, self.rebaseline_reason),
            space_heating_kwh=self.space_heating_kwh,
            water_heating_kwh=self.water_heating_kwh,
            bill=self._read_bill(),
        )

    def _read_lodged(self) -> Optional[Performance]:
        """The Lodged half, or None for a predicted Property whose four lodged_*
        columns are NULL (ADR-0004 amendment, #1361). They are written as a unit
        (`from_domain`), so lodged_sap_score is the discriminator and the other
        three are non-null alongside it."""
        if self.lodged_sap_score is None:
            return None
        assert self.lodged_epc_band is not None
        assert self.lodged_co2_emissions_t_per_yr is not None
        assert self.lodged_primary_energy_intensity_kwh_per_m2_yr is not None
        return Performance(
            sap_score=self.lodged_sap_score,
            epc_band=Epc(self.lodged_epc_band),
            co2_emissions=self.lodged_co2_emissions_t_per_yr,
            primary_energy_intensity=self.lodged_primary_energy_intensity_kwh_per_m2_yr,
        )

    def _read_bill(self) -> Optional[Bill]:
        """Reconstruct the Bill from the ``bill_*`` columns. The total is the
        not-None discriminator: a persisted bill always sets it, so its absence
        means no calculator ran and the bill was None. A section is rebuilt only
        when its kWh column is not None (paired with its cost)."""
        if self.total_annual_bill_gbp is None:
            return None
        sections: dict[BillSection, BillSectionCost] = {}
        for section, stem in _SECTION_COLUMN_STEM.items():
            kwh = cast(Optional[float], getattr(self, f"{stem}_kwh"))
            if kwh is None:
                continue
            cost_gbp = cast(float, getattr(self, f"{stem}_cost_gbp"))
            sections[section] = BillSectionCost(kwh=kwh, cost_gbp=cost_gbp)
        return Bill(
            sections=sections,
            standing_charges_gbp=cast(float, self.standing_charges_gbp),
            seg_credit_gbp=cast(float, self.seg_credit_gbp),
            total_gbp=self.total_annual_bill_gbp,
        )