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.

The DB engine is module-scoped (ADR-0012). Architecturally each invocation uses
one DB connection at a time: 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 whose
overrides resolve on its own session, so no two Sessions ever overlap. The engine
uses ``NullPool`` rather than a fixed pool so that target is a graceful ceiling,
not a hard one: a fresh connection is opened per checkout and closed on return,
so there is no pool slot to exhaust — any future accidental overlap opens a
transient second connection instead of dead-locking the Lambda.
"""

from __future__ import annotations

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 sqlalchemy.pool import NullPool
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))
        # Architecturally 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 — and the Unit of
        # Work resolves overrides on its own session — so no two Sessions overlap
        # and a single connection suffices. 32 concurrent containers × 1 = 32
        # against RDS.
        #
        # NullPool, not a fixed pool, enforces that as a *graceful* ceiling rather
        # than a hard one: each checkout opens a fresh connection and closes it on
        # return, so there is no pool slot to exhaust. If a future code path ever
        # holds two Sessions at once it opens a second connection for that instant
        # instead of dead-locking on a 1-slot pool and failing the whole
        # invocation (the "QueuePool limit of size 1 overflow 0 reached" timeout).
        # The design target stays one connection; NullPool just keeps the Lambda
        # running if we ever regress it.
        _engine = make_engine(config, poolclass=NullPool)
    return _engine


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

    Its repositories commit on every ``save``/``create``, releasing the
    connection between bookkeeping calls, so it holds none while the wrapped
    handler body runs. Combined with the read-then-write handler structure, 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()