Model/applications/modelling_e2e/handler.py

"""SQS-triggered Lambda: fetch EPC → run modelling → persist plan.

One SQS message = one property. The handler reads ``property_id``,
``portfolio_id``, ``scenario_id``, and ``no_solar`` from the message body,
fetches the property's EPC from the gov API, runs the full modelling pipeline
(SAP10 → optimiser) via ``harness.console.run_modelling``, and persists the
resulting Plan via ``PlanPostgresRepository.save()``.

``secondary_heating_removal`` is excluded unconditionally: the live ``material``
catalogue does not yet carry this measure type, causing a crash during catalogue
reads for properties with a lodged secondary heater.

DB engine is module-scoped so the connection pool is reused across warm
invocations (ADR-0012).
"""

from __future__ import annotations

import io
import os
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 EpcPropertyData
from domain.geospatial.planning_restrictions import PlanningRestrictions
from domain.geospatial.spatial_reference import SpatialReference
from domain.modelling.measure_type import MeasureType
from domain.property.property import Property, PropertyIdentity
from domain.tasks.tasks import Source
from harness.console import run_modelling
from infrastructure.epc_client.epc_client_service import EpcClientService
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.modelling_e2e_trigger_body import (
    ModellingE2ETriggerBody,
)
from repositories.geospatial.geospatial_s3_repository import (
    GeospatialS3Repository,
    ParquetReader,
)
from repositories.epc.epc_postgres_repository import EpcPostgresRepository
from repositories.plan.plan_postgres_repository import PlanPostgresRepository
from repositories.product.product_postgres_repository import ProductPostgresRepository
from repositories.property.landlord_override_overlays import overlays_from
from repositories.property.property_overrides_postgres_reader import (
    PropertyOverridesPostgresReader,
)
from repositories.property.property_postgres_repository import (
    PropertyPostgresRepository,
)
from repositories.scenario.scenario_postgres_repository import (
    ScenarioPostgresRepository,
)
from utilities.aws_lambda.task_handler import task_handler
from utilities.logger import setup_logger

_engine: Optional[Engine] = None

logger = setup_logger()


def _get_engine() -> Engine:
    global _engine
    if _engine is None:
        _engine = make_engine(PostgresConfig.from_env(dict(os.environ)))
    return _engine


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


@task_handler(task_source="modelling_e2e", source=Source.PROPERTY)
def handler(body: dict[str, Any], context: Any) -> None:
    trigger = ModellingE2ETriggerBody.model_validate(body)
    property_id = trigger.property_id
    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={property_id} 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:
        row = conn.execute(
            text("SELECT uprn FROM property WHERE id = :pid"),
            {"pid": property_id},
        ).one()
    uprn = int(row[0])
    logger.info(f"resolved uprn={uprn}")

    epc: Optional[EpcPropertyData] = epc_client.get_by_uprn(uprn)
    if epc is None:
        raise ValueError(f"no EPC found for UPRN {uprn} (property {property_id})")
    logger.info(f"fetched EPC (energy_rating_current={epc.energy_rating_current})")

    overrides_reader = PropertyOverridesPostgresReader(lambda: Session(engine))
    overlaid = Property(
        identity=PropertyIdentity(
            portfolio_id=portfolio_id, postcode="", address="", uprn=uprn
        ),
        epc=epc,
        landlord_overrides=overlays_from(overrides_reader.overrides_for(property_id)),
    )
    effective_epc = overlaid.effective_epc

    spatial = _spatial_for(geospatial, uprn)
    restrictions = (
        spatial.restrictions if spatial is not None else PlanningRestrictions()
    )
    logger.info(f"spatial={'found' if spatial is not None else 'not found'}")

    if no_solar:
        solar_insights = None
        logger.info("solar skipped (no_solar=True)")
    else:
        solar_insights = _solar_insights_for(solar_client, spatial)
        logger.info(f"solar={'found' if solar_insights is not None else 'not found'}")

    with Session(engine) as session:
        scenario = ScenarioPostgresRepository(session).get_many([scenario_id])[0]
        logger.info(f"loaded scenario goal={scenario.goal!r} goal_value={scenario.goal_value!r}")
        products = ProductPostgresRepository(session)

        # secondary_heating_removal is absent from the live material.type enum;
        # exclude it unconditionally until the catalogue gap is resolved.
        considered: Optional[frozenset[MeasureType]] = frozenset(MeasureType) - {
            MeasureType.SECONDARY_HEATING_REMOVAL
        }

        logger.info("running modelling pipeline")
        plan = run_modelling(
            effective_epc,
            planning_restrictions=restrictions,
            solar_insights=solar_insights,
            considered_measures=considered,
            products=products,
            scenario=scenario,
            print_table=False,
        )
        logger.info(
            f"modelling complete: SAP {plan.baseline.sap_continuous:.1f}→"
            f"{plan.post_sap_continuous:.1f} measures={len(plan.measures)} "
            f"cost=£{plan.cost_of_works:,.0f}"
        )

        if dry_run:
            measure_types = ", ".join(m.measure_type for m in plan.measures) or "none"
            logger.info(f"[dry_run] measures=[{measure_types}] — skipping DB write")
            return
        EpcPostgresRepository(session).save(
            epc, property_id=property_id, portfolio_id=portfolio_id
        )
        PlanPostgresRepository(session).save(
            plan,
            property_id=property_id,
            scenario_id=scenario_id,
            portfolio_id=portfolio_id,
            is_default=scenario.is_default,
        )
        PropertyPostgresRepository(session).mark_modelled(
            property_id, has_recommendations=bool(plan.measures)
        )
        session.commit()
        logger.info("plan saved")