Model/applications/modelling_e2e/handler.py

"""SQS-triggered Lambda: fetch EPC (or predict) → run modelling → persist plan.

One SQS message = one batch of properties sharing a portfolio, scenario, and
(by caller convention) postcode. The handler reads ``property_ids``,
``portfolio_id``, ``scenario_id``, ``no_solar``, and ``dry_run`` from the
message body, fetches or predicts each property's EPC, runs the full modelling
pipeline (SAP10 → optimiser) via ``harness.console.run_modelling``, and
persists the resulting Plan via ``PostgresUnitOfWork`` in one atomic transaction
per property.

When no lodged EPC is found, EPC Prediction (Path 3, ADR-0031) synthesises one
from the postcode cohort. ``_cohort_cache`` is module-level so warm Lambda
containers re-processing the same postcode avoid redundant fetches.

All Measure Types are considered: pricing goes through
``catalogue_snapshot_with_off_catalogue_overrides`` so the measures the live
``material`` catalogue cannot supply (``secondary_heating_removal``, the glazing
and heating gaps) are priced from the committed off-catalogue overlay instead of
crashing.

DB engine is module-scoped so the connection pool is reused across warm
invocations (ADR-0012). The pool holds a single connection (``pool_size=1``): the
handler reads everything up front — overrides, Scenario, a catalogue snapshot, and
stored Solar — through one short-lived read Session, closes it, then writes each
Property in a sequential Unit of Work, so the read and write Sessions never
overlap. The orchestrator shares the same engine and releases its connection
between bookkeeping commits, so one invocation uses one DB connection at a time.
"""

from __future__ import annotations

import dataclasses
import io
import os
from collections.abc import Callable, Generator
from contextlib import contextmanager
from typing import Any, Optional, cast

import boto3
import pandas as pd  # pyright: ignore[reportMissingTypeStubs]
from sqlalchemy import Engine, text
from sqlmodel import Session

from datatypes.epc.domain.epc_property_data import (
    BuildingPartIdentifier,
    EpcPropertyData,
)
from domain.epc_prediction.comparable_properties import (
    ComparableProperty,
    select_comparables,
)
from domain.epc_prediction.epc_prediction import EpcPrediction
from domain.epc_prediction.prediction_target import (
    PredictionTarget,
    build_prediction_target,
)
from domain.geospatial.coordinates import Coordinates
from domain.geospatial.planning_restrictions import PlanningRestrictions
from domain.geospatial.spatial_reference import SpatialReference
from domain.property.property import Property, PropertyIdentity
from domain.property_baseline.calculator_rebaseliner import CalculatorRebaseliner
from domain.sap10_calculator.calculator import Sap10Calculator
from domain.tasks.tasks import Source
from harness.console import run_modelling
from orchestration.task_orchestrator import TaskOrchestrator
from orchestration.property_baseline_orchestrator import (
    PropertyBaselineOrchestrator,
)
from infrastructure.epc_client.epc_client_service import EpcClientService
from infrastructure.postcodes_io.postcodes_io_client import PostcodesIoClient
from infrastructure.postgres.config import PostgresConfig
from infrastructure.postgres.engine import make_engine
from infrastructure.solar.google_solar_api_client import (
    BuildingInsightsNotFoundError,
    GoogleSolarApiClient,
)
from applications.modelling_e2e.errors import (
    DegeneratePredictionError,
    NoSameTypeComparablesError,
    UnresolvedPropertyTypeError,
)
from applications.modelling_e2e.modelling_e2e_trigger_body import (
    ModellingE2ETriggerBody,
)
from repositories.comparable_properties.epc_comparable_properties_repository import (
    EpcComparablePropertiesRepository,
    SkippedCohortCert,
)
from repositories.geospatial.geospatial_s3_repository import (
    GeospatialS3Repository,
    ParquetReader,
)
from repositories.fuel_rates.fuel_rates_static_file_repository import (
    FuelRatesStaticFileRepository,
)
from repositories.postgres_unit_of_work import PostgresUnitOfWork
from repositories.product.composite_product_repository import (
    catalogue_snapshot_with_off_catalogue_overrides,
)
from repositories.property.in_memory_property_overrides_reader import (
    InMemoryPropertyOverridesReader,
)
from repositories.property.landlord_override_overlays import overlays_from
from repositories.property.override_backed_prediction_attributes_reader import (
    OverrideBackedPredictionAttributesReader,
)
from repositories.property.property_overrides_postgres_reader import (
    PropertyOverridesPostgresReader,
)
from repositories.property.property_overrides_reader import (
    ResolvedPropertyOverrides,
)
from repositories.scenario.scenario_postgres_repository import (
    ScenarioPostgresRepository,
)
from repositories.solar.solar_postgres_repository import SolarPostgresRepository
from repositories.tasks.subtask_postgres_repository import (
    SubTaskPostgresRepository,
)
from repositories.tasks.task_postgres_repository import TaskPostgresRepository
from utilities.aws_lambda.task_handler import task_handler
from uuid import UUID
from utilities.logger import setup_logger

_engine: Optional[Engine] = None
_cohort_cache: dict[str, list[ComparableProperty]] = {}
# Broadened (nearby-postcode) cohorts, keyed by (seed postcode, target property
# type): the early-stop walk depends on the type it is filling for, so two types
# in the same postcode must not share a cached result.
_nearby_cohort_cache: dict[tuple[str, str], list[ComparableProperty]] = {}

logger = setup_logger()


def _get_engine() -> Engine:
    global _engine
    if _engine is None:
        config = PostgresConfig.from_env(dict(os.environ))
        # One connection per invocation: the handler reads everything up front
        # through one short-lived read Session, closes it, then writes each
        # Property in a sequential Unit of Work — so the read and write Sessions
        # never overlap and a single pooled connection suffices. The orchestrator
        # shares this engine (see ``_shared_engine_orchestrator``) and releases
        # its connection between bookkeeping commits, so it holds none during the
        # work. 32 concurrent containers × 1 connection = 32 against RDS.
        _engine = make_engine(dataclasses.replace(config, pool_size=1, max_overflow=0))
    return _engine


@contextmanager
def _shared_engine_orchestrator() -> Generator[TaskOrchestrator, None, None]:
    """A ``TaskOrchestrator`` on the same module-scoped pooled engine as the
    modelling work — not a separate per-invocation NullPool engine.

    Its repositories commit on every ``save``/``create``, releasing the pooled
    connection between bookkeeping calls, so it holds none while the wrapped
    handler body runs. Combined with the read-then-write handler structure and
    ``pool_size=1``, the whole invocation uses one DB connection at a time."""
    engine = _get_engine()
    with Session(engine) as session:
        yield TaskOrchestrator(
            task_repo=TaskPostgresRepository(session=session),
            subtask_repo=SubTaskPostgresRepository(session=session),
        )


def _s3_parquet_reader() -> ParquetReader:
    bucket = os.environ["DATA_BUCKET"]

    def read(key: str) -> pd.DataFrame:
        s3: Any = cast(
            Any, boto3.client("s3")
        )  # pyright: ignore[reportUnknownMemberType]
        raw = cast(bytes, s3.get_object(Bucket=bucket, Key=key)["Body"].read())
        return pd.read_parquet(io.BytesIO(raw))  # type: ignore[return-value]

    return read


def _spatial_for(
    geospatial: GeospatialS3Repository, uprn: int
) -> Optional[SpatialReference]:
    try:
        return geospatial.spatial_for(uprn)
    except Exception:  # noqa: BLE001
        return None


def _solar_insights_for(
    solar_client: GoogleSolarApiClient, spatial: Optional[SpatialReference]
) -> Optional[dict[str, Any]]:
    if spatial is None or spatial.coordinates is None:
        return None
    try:
        return solar_client.get_building_insights(
            spatial.coordinates.longitude, spatial.coordinates.latitude
        )
    except BuildingInsightsNotFoundError:
        return None


def _dedupe_skipped(
    skipped: list[SkippedCohortCert],
) -> list[SkippedCohortCert]:
    """First occurrence of each skipped cert number (the same cert can appear in
    more than one postcode cohort across a batch)."""
    seen: set[str] = set()
    unique: list[SkippedCohortCert] = []
    for cert in skipped:
        if cert.certificate_number not in seen:
            seen.add(cert.certificate_number)
            unique.append(cert)
    return unique


def _predict_epc(
    *,
    property_id: int,
    uprn: int,
    postcode: str,
    portfolio_id: int,
    attributes_reader: OverrideBackedPredictionAttributesReader,
    coordinates: Optional[Coordinates],
    cohort_for: Callable[[str], list[ComparableProperty]],
    broaden: Callable[[PredictionTarget], list[ComparableProperty]],
    predictor: EpcPrediction,
) -> EpcPropertyData:
    """Synthesise an EpcPropertyData for an EPC-less property from its postcode
    cohort (EPC Prediction Path 3, ADR-0031).

    When the property's own postcode holds no same-type comparables (a sparse
    postcode — e.g. the only flat among houses), the cohort is broadened to the
    real unit postcodes physically nearest it (``broaden``) before giving up.

    Raises a specific ``PropertyNotModellableError`` subclass — naming the cause
    and carrying the property's identity — when it cannot predict: property_type
    unresolved, an empty same-type cohort, or a degenerate (no MAIN part)
    prediction. The per-property handler records ``str(exc)`` in the SubTask
    output, so the cause is debuggable from the output alone.
    """
    attributes = attributes_reader.attributes_for(property_id)
    identity = PropertyIdentity(
        portfolio_id=portfolio_id, postcode=postcode, address="", uprn=uprn
    )
    target = build_prediction_target(identity, coordinates, attributes)
    if target is None:
        raise UnresolvedPropertyTypeError(
            property_id=property_id,
            uprn=uprn,
            postcode=postcode,
            portfolio_id=portfolio_id,
            property_type=attributes.property_type,
            built_form=attributes.built_form,
        )
    comparables = select_comparables(target, cohort_for(target.postcode))
    broadened = False
    if not comparables.members:
        broadened = True
        comparables = select_comparables(target, broaden(target))
    if not comparables.members:
        raise NoSameTypeComparablesError(
            property_id=property_id,
            uprn=uprn,
            postcode=postcode,
            portfolio_id=portfolio_id,
            property_type=target.property_type,
            broadened=broadened,
        )
    predicted = predictor.predict(target, comparables)
    if not any(
        part.identifier is BuildingPartIdentifier.MAIN
        for part in predicted.sap_building_parts
    ):
        raise DegeneratePredictionError(
            property_id=property_id,
            uprn=uprn,
            postcode=postcode,
            portfolio_id=portfolio_id,
            property_type=target.property_type,
            cohort_size=len(comparables.members),
        )
    return predicted


@task_handler(
    task_source="modelling_e2e",
    source=Source.PROPERTY,
    orchestrator_cm=_shared_engine_orchestrator,
    pass_task_orchestrator=True,
)
def handler(body: dict[str, Any], context: Any, orchestrator: TaskOrchestrator, task_id: UUID) -> None:
    trigger = ModellingE2ETriggerBody.model_validate(body)
    property_ids = trigger.property_ids
    portfolio_id = trigger.portfolio_id
    scenario_id = trigger.scenario_id
    no_solar = trigger.no_solar
    dry_run = trigger.dry_run

    logger.info(
        f"start property_ids={property_ids} portfolio={portfolio_id} "
        f"scenario={scenario_id} no_solar={no_solar} dry_run={dry_run}"
    )

    engine = _get_engine()
    epc_client = EpcClientService(os.environ["OPEN_EPC_API_TOKEN"])
    geospatial = GeospatialS3Repository(_s3_parquet_reader())
    solar_client = GoogleSolarApiClient(os.environ["GOOGLE_SOLAR_API_KEY"])

    with engine.connect() as conn:
        uprn_rows = conn.execute(
            text("SELECT id, uprn FROM property WHERE id = ANY(:ids)"),
            {"ids": property_ids},
        ).fetchall()
        postcode_rows = conn.execute(
            text("SELECT id, postcode FROM property WHERE id = ANY(:ids)"),
            {"ids": property_ids},
        ).fetchall()

    uprns: dict[int, int] = {int(row[0]): int(row[1]) for row in uprn_rows}
    postcodes: dict[int, str] = {int(row[0]): (row[1] or "") for row in postcode_rows}

    # Pre-fetch every Property's overrides up front (each call opens and closes
    # its own short read Session) and serve them from memory through the loop, so
    # no override read Session is held open alongside a write Unit of Work.
    overrides_postgres_reader = PropertyOverridesPostgresReader(lambda: Session(engine))
    overrides_by_pid: dict[int, ResolvedPropertyOverrides] = {
        pid: overrides_postgres_reader.overrides_for(pid) for pid in property_ids
    }
    overrides_reader = InMemoryPropertyOverridesReader(overrides_by_pid)
    prediction_attrs_reader = OverrideBackedPredictionAttributesReader(overrides_reader)
    comparables_repo = EpcComparablePropertiesRepository(
        epc_client, geospatial, nearby_postcodes=PostcodesIoClient()
    )
    predictor = EpcPrediction()

    def _get_cohort(postcode: str) -> list[ComparableProperty]:
        if postcode not in _cohort_cache:
            _cohort_cache[postcode] = (
                comparables_repo.candidates_for(postcode) if postcode else []
            )
        return _cohort_cache[postcode]

    def _broaden(target: PredictionTarget) -> list[ComparableProperty]:
        """The nearby-postcode cohort for a gated-out target — the real unit
        postcodes nearest it, walked until enough same-type comparables surface
        (ADR-0034). Memoised per (postcode, property_type) so co-located
        same-type misses share one walk."""
        key = (target.postcode, target.property_type)
        if key not in _nearby_cohort_cache:
            _nearby_cohort_cache[key] = (
                comparables_repo.candidates_near(
                    target.postcode,
                    target.coordinates,
                    enough=lambda c: c.epc.property_type == target.property_type,
                )
                if target.postcode
                else []
            )
        return _nearby_cohort_cache[key]

    # Re-establishes each lodged Property's Baseline Performance from the just-
    # persisted EPC (one UoW per property, committed after the Plan's). Predicted
    # Properties have no lodged figures, so they get no baseline (mirrors the e2e
    # runner and the ara_first_run Baseline stage).
    baseline_orchestrator = PropertyBaselineOrchestrator(
        unit_of_work=lambda: PostgresUnitOfWork(lambda: Session(engine)),
        rebaseliner=CalculatorRebaseliner(Sap10Calculator()),
        fuel_rates=FuelRatesStaticFileRepository(),
    )

    read_session = Session(engine)
    try:
        # Read everything the write loop needs up front: the Scenario, an in-memory
        # snapshot of the catalogue (priced after the Session closes), and each
        # UPRN's stored Solar insights. Then close the read Session immediately so
        # its pooled connection is free before the loop — each Property's write
        # Unit of Work reuses that single connection rather than opening a second
        # alongside a held-open read Session. (The ``finally`` is the safety net.)
        scenario = ScenarioPostgresRepository(read_session).get_many([scenario_id])[0]
        products = catalogue_snapshot_with_off_catalogue_overrides(read_session)
        stored_solar: dict[int, Optional[dict[str, Any]]] = (
            {}
            if no_solar
            else {
                uprn: SolarPostgresRepository(read_session).get(uprn)
                for uprn in set(uprns.values())
            }
        )
        read_session.close()

        for property_id in property_ids:
            child = orchestrator.create_child_subtask(
                task_id, inputs={"property_id": property_id}
            )

            def _work(pid: int = property_id) -> None:
                uprn = uprns[pid]
                postcode = postcodes.get(pid, "")
                logger.info(f"property={pid} uprn={uprn} postcode={postcode!r}")

                spatial = _spatial_for(geospatial, uprn)
                restrictions = (
                    spatial.restrictions
                    if spatial is not None
                    else PlanningRestrictions()
                )
                coordinates: Optional[Coordinates] = (
                    spatial.coordinates if spatial is not None else None
                )

                epc: Optional[EpcPropertyData] = epc_client.get_by_uprn(uprn)
                overrides = overlays_from(overrides_reader.overrides_for(pid))
                predicted_epc: Optional[EpcPropertyData] = None

                if epc is not None:
                    logger.info(f"property={pid} lodged EPC found")
                    effective_epc = Property(
                        identity=PropertyIdentity(
                            portfolio_id=portfolio_id,
                            postcode=postcode,
                            address="",
                            uprn=uprn,
                        ),
                        epc=epc,
                        landlord_overrides=overrides,
                    ).effective_epc
                else:
                    logger.info(
                        f"property={pid} no lodged EPC — attempting prediction"
                    )
                    predicted_epc = _predict_epc(
                        property_id=pid,
                        uprn=uprn,
                        postcode=postcode,
                        portfolio_id=portfolio_id,
                        attributes_reader=prediction_attrs_reader,
                        coordinates=coordinates,
                        cohort_for=_get_cohort,
                        broaden=_broaden,
                        predictor=predictor,
                    )
                    effective_epc = Property(
                        identity=PropertyIdentity(
                            portfolio_id=portfolio_id,
                            postcode=postcode,
                            address="",
                            uprn=uprn,
                        ),
                        epc=None,
                        predicted_epc=predicted_epc,
                        landlord_overrides=overrides,
                    ).effective_epc

                solar_insights: Optional[dict[str, Any]]
                solar_was_fetched = False
                if no_solar:
                    solar_insights = None
                else:
                    solar_insights = stored_solar.get(uprn)
                    if solar_insights is None:
                        solar_insights = _solar_insights_for(solar_client, spatial)
                        solar_was_fetched = solar_insights is not None

                plan = run_modelling(
                    effective_epc,
                    planning_restrictions=restrictions,
                    solar_insights=solar_insights,
                    considered_measures=None,
                    products=products,
                    scenario=scenario,
                    print_table=False,
                )
                logger.info(
                    f"property={pid} modelling complete "
                    f"measures={len(plan.measures)}"
                )

                if dry_run:
                    measure_types = (
                        ", ".join(m.measure_type for m in plan.measures) or "none"
                    )
                    logger.info(
                        f"[dry_run] property={pid} "
                        f"measures=[{measure_types}] — skipping DB write"
                    )
                    return

                with PostgresUnitOfWork(lambda: Session(engine)) as uow:
                    if epc is not None:
                        uow.epc.save(
                            epc, property_id=pid, portfolio_id=portfolio_id
                        )
                    elif predicted_epc is not None:
                        # Persist the synthesised EPC in the predicted slot (ADR-0031),
                        # so the Baseline stage can re-hydrate it and downstream sees
                        # the picture the Plan was modelled from.
                        uow.epc.save(
                            predicted_epc,
                            property_id=pid,
                            portfolio_id=portfolio_id,
                            source="predicted",
                        )
                    if spatial is not None:
                        uow.spatial.save(uprn, spatial)
                    if (
                        solar_was_fetched
                        and solar_insights is not None
                        and spatial is not None
                        and spatial.coordinates is not None
                    ):
                        uow.solar.save(
                            uprn,
                            longitude=spatial.coordinates.longitude,
                            latitude=spatial.coordinates.latitude,
                            insights=solar_insights,
                        )
                    uow.plan.save(
                        plan,
                        property_id=pid,
                        scenario_id=scenario_id,
                        portfolio_id=portfolio_id,
                        is_default=scenario.is_default,
                    )
                    uow.property.mark_modelled(
                        pid, has_recommendations=bool(plan.measures)
                    )
                    uow.commit()
                logger.info(f"property={pid} plan saved")

                baseline_orchestrator.run([pid])
                logger.info(f"property={pid} baseline saved")

            try:
                orchestrator.run_subtask(child.id, work=_work)
            except Exception:  # noqa: BLE001
                pass

        # Cohort certs the mapper could not consume were skipped (not aborted on)
        # so prediction could proceed; surface them — with cert numbers — in the
        # subtask outputs so the mapper gaps can be closed later.
        skipped_certs: list[dict[str, str]] = [
            {"certificate_number": s.certificate_number, "error": s.error}
            for s in _dedupe_skipped(comparables_repo.skipped)
        ]
        if skipped_certs:
            logger.info(
                f"skipped {len(skipped_certs)} unmappable cohort cert(s): "
                f"{[s['certificate_number'] for s in skipped_certs]}"
            )

    finally:
        read_session.close()