Model/orchestration/ingestion_orchestrator.py

from __future__ import annotations

from collections.abc import Callable
from dataclasses import dataclass
from typing import Any, Optional, Protocol

from datatypes.epc.domain.epc_property_data import EpcPropertyData
from datatypes.epc.domain.historic_epc import HistoricEpc
from domain.epc_prediction.comparable_properties import (
    ComparableProperty,
    select_comparables,
)
from domain.epc_prediction.epc_prediction import EpcPrediction
from domain.epc_prediction.historic_conditioning import (
    HistoricConditioning,
    attributes_with_historic_fallback,
    conditioning_from_historic,
    target_with_conditioning,
)
from domain.epc_prediction.prediction_target import (
    PredictionTargetAttributes,
    build_prediction_target,
)
from domain.geospatial.coordinates import Coordinates
from domain.geospatial.spatial_reference import SpatialReference
from domain.property.property import PropertyIdentity
from repositories.epc.epc_repository import EpcSource
from repositories.geospatial.geospatial_repository import GeospatialRepository
from repositories.unit_of_work import UnitOfWork


class EpcFetcher(Protocol):
    """The slice of the New-EPC-API client Ingestion needs (e.g. EpcClientService)."""

    def get_by_uprn(self, uprn: int) -> Optional[EpcPropertyData]: ...


class ComparablesRepo(Protocol):
    """The cohort source for EPC Prediction (e.g. EpcComparablePropertiesRepository)."""

    def candidates_for(self, postcode: str) -> list[ComparableProperty]: ...


class PredictionAttributesReader(Protocol):
    """Resolves an EPC-less Property's prediction attributes from Landlord
    Overrides (e.g. the property_overrides read adapter)."""

    def attributes_for(self, property_id: int) -> PredictionTargetAttributes: ...


class HistoricEpcReader(Protocol):
    """The slice of the Historic-EPC resolver Ingestion needs: the exact-UPRN
    lookup whose stable attributes condition prediction (ADR-0054)."""

    def record_for_uprn(self, uprn: str, postcode: str) -> Optional[HistoricEpc]: ...


class SolarFetcher(Protocol):
    """The slice of the Google Solar client Ingestion needs (e.g. GoogleSolarApiClient)."""

    def get_building_insights(
        self, longitude: float, latitude: float
    ) -> dict[str, Any]: ...


@dataclass
class _Prep:
    """A property's transactional inputs read in the unit phase, before external
    IO: its identity (postcode + uprn) and, when the predictor is wired, its
    resolved prediction attributes (so the no-unit fetch phase can predict)."""

    property_id: int
    identity: PropertyIdentity
    attributes: Optional[PredictionTargetAttributes]


@dataclass
class _Fetched:
    """One property's externally-fetched source data, awaiting the write phase."""

    property_id: int
    uprn: int
    epc: Optional[EpcPropertyData]
    predicted_epc: Optional[EpcPropertyData]
    solar_insights: Optional[dict[str, Any]]
    spatial: Optional[SpatialReference]
    # "expired" when the prediction was conditioned by an expired Historic EPC
    # (ADR-0054); plain "predicted" otherwise.
    predicted_source: EpcSource = "predicted"


class IngestionOrchestrator:
    """Stage 1: acquire a batch's external source data and persist it.

    Runs in two phases so a DB connection is never held during external IO
    (ADR-0012): **fetch** the whole batch — read each UPRN, fetch its EPC,
    resolve its spatial reference (coordinates + planning protections) from the
    Geospatial reference Repo, thread the coordinates into the Solar fetcher —
    with *no unit open*; then **write** the batch in one Unit of Work and commit
    once. Fetchers never call each other (ADR-0011); the orchestrator threads
    the coordinate. The coordinates drive the Solar fetch transiently; the whole
    spatial reference is cached per-UPRN in the transactional store so Modelling
    reads the planning protections back off the Property (ADR-0020).

    The geospatial repo reads S3 reference data, not the transactional store, so
    it is injected separately rather than taken from the unit.
    """

    def __init__(
        self,
        *,
        unit_of_work: Callable[[], UnitOfWork],
        epc_fetcher: EpcFetcher,
        geospatial_repo: GeospatialRepository,
        solar_fetcher: SolarFetcher,
        comparables_repo: Optional[ComparablesRepo] = None,
        prediction_attributes_reader: Optional[PredictionAttributesReader] = None,
        epc_prediction: Optional[EpcPrediction] = None,
        historic_epc_reader: Optional[HistoricEpcReader] = None,
    ) -> None:
        self._unit_of_work = unit_of_work
        self._epc_fetcher = epc_fetcher
        self._geospatial_repo = geospatial_repo
        self._solar_fetcher = solar_fetcher
        # EPC Prediction gap-fill (ADR-0031): when all three are wired, an EPC-less
        # Property is predicted from its postcode cohort and persisted to the
        # predicted slot. When any is absent, prediction is simply off and
        # ingestion behaves exactly as before.
        self._comparables_repo = comparables_repo
        self._prediction_attributes_reader = prediction_attributes_reader
        self._epc_prediction = epc_prediction
        self._historic_epc_reader = historic_epc_reader

    def run(self, property_ids: list[int]) -> None:
        preps = self._prepare(property_ids)
        fetched = [self._fetch(prep) for prep in preps]
        self._persist(fetched)

    def _prepare(self, property_ids: list[int]) -> list[_Prep]:
        # A short read unit; properties with no UPRN (e.g. landlord_property_id
        # only) are skipped — a later Site-Notes path covers them. Prediction
        # attributes (Landlord Overrides) are resolved here, in-unit, so the
        # no-unit fetch phase holds everything it needs to predict.
        with self._unit_of_work() as uow:
            properties = uow.property.get_many(property_ids)
            preps: list[_Prep] = []
            for property_id, prop in zip(property_ids, properties, strict=True):
                if prop.identity.uprn is None:
                    continue
                attributes = (
                    self._prediction_attributes_reader.attributes_for(property_id)
                    if self._prediction_attributes_reader is not None
                    else None
                )
                preps.append(_Prep(property_id, prop.identity, attributes))
            return preps

    def _fetch(self, prep: _Prep) -> _Fetched:
        # No unit open here — this is the external-IO phase. One spatial
        # reference lookup yields the coordinates (which drive the Solar fetch)
        # and the planning protections (cached for Modelling, ADR-0020).
        uprn = prep.identity.uprn
        assert uprn is not None  # _prepare drops UPRN-less properties
        epc = self._epc_fetcher.get_by_uprn(uprn)
        solar_insights: Optional[dict[str, Any]] = None
        spatial: Optional[SpatialReference] = self._geospatial_repo.spatial_for(uprn)
        coordinates = spatial.coordinates if spatial is not None else None
        if coordinates is not None:
            solar_insights = self._solar_fetcher.get_building_insights(
                coordinates.longitude, coordinates.latitude
            )
        predicted_epc: Optional[EpcPropertyData] = None
        conditioning: Optional[HistoricConditioning] = None
        if epc is None:
            conditioning = self._historic_conditioning(uprn, prep.identity.postcode)
            predicted_epc = self._predict(
                prep.identity, coordinates, prep.attributes, conditioning
            )
        predicted_source = (
            "expired"
            if predicted_epc is not None and conditioning is not None
            else "predicted"
        )
        return _Fetched(
            prep.property_id,
            uprn,
            epc,
            predicted_epc,
            solar_insights,
            spatial,
            predicted_source,
        )

    def _historic_conditioning(
        self, uprn: int, postcode: str
    ) -> Optional[HistoricConditioning]:
        """The expired Historic EPC's stable attributes for this exact UPRN, or
        None when the reader is unwired or the backup holds no row (ADR-0054)."""
        if self._historic_epc_reader is None:
            return None
        record = self._historic_epc_reader.record_for_uprn(str(uprn), postcode)
        return conditioning_from_historic(record) if record is not None else None

    def _predict(
        self,
        identity: PropertyIdentity,
        coordinates: Optional[Coordinates],
        attributes: Optional[PredictionTargetAttributes],
        conditioning: Optional[HistoricConditioning] = None,
    ) -> Optional[EpcPropertyData]:
        """Synthesise the EPC-less Property's picture from its postcode cohort, or
        None when the predictor is unwired, the Property is gated out (unknown
        property type), or no comparables survive selection (ADR-0031). An
        expired Historic EPC's stable attributes fill override gaps and condition
        the cohort; Landlord Overrides always win where both speak (ADR-0054)."""
        if (
            self._comparables_repo is None
            or self._epc_prediction is None
            or (attributes is None and conditioning is None)
        ):
            return None
        if conditioning is not None:
            attributes = attributes_with_historic_fallback(attributes, conditioning)
        assert attributes is not None  # one of the two branches above supplied it
        target = build_prediction_target(identity, coordinates, attributes)
        if target is None:
            return None
        if conditioning is not None:
            target = target_with_conditioning(target, conditioning)
        candidates = self._comparables_repo.candidates_for(identity.postcode)
        comparables = select_comparables(target, candidates)
        if not comparables.members:
            return None
        return self._epc_prediction.predict(target, comparables)

    def _persist(self, fetched: list[_Fetched]) -> None:
        with self._unit_of_work() as uow:
            for item in fetched:
                if item.epc is not None:
                    uow.epc.save(item.epc, property_id=item.property_id)
                elif item.predicted_epc is not None:
                    uow.epc.save(
                        item.predicted_epc,
                        property_id=item.property_id,
                        source=item.predicted_source,
                    )
                # The live `solar` table is keyed by UPRN and needs the fetch's
                # coordinates; insights are only set when those coordinates were
                # resolved, so spatial.coordinates is non-None alongside them.
                if (
                    item.solar_insights is not None
                    and item.spatial is not None
                    and item.spatial.coordinates is not None
                ):
                    uow.solar.save(
                        item.uprn,
                        longitude=item.spatial.coordinates.longitude,
                        latitude=item.spatial.coordinates.latitude,
                        insights=item.solar_insights,
                    )
                if item.spatial is not None:
                    uow.spatial.save(item.uprn, item.spatial)
            uow.commit()