Model/scripts/run_first_run_e2e.py

"""Run the **full** ``AraFirstRunPipeline`` (Ingestion → Baseline → Modelling)
end-to-end against the real database, locally.

This is the production pipeline the ``ara_first_run`` Lambda runs, driven from a
shell instead of an SQS event. The Lambda ``handler`` itself cannot run locally —
``applications/ara_first_run/handler.py::_source_clients_from_env`` deliberately
raises until the deploy/Terraform wiring lands (#1136). So this script composes
the same pipeline directly via the existing ``build_first_run_pipeline`` seam,
supplying the three source clients that ``run_modelling_e2e`` already proves out
(EPC API, geospatial S3, Google Solar), then calls ``dispatch_first_run``.

How it differs from ``run_modelling_e2e``:
  * It runs the **real Ingestion stage** — fetches each Property's EPC by UPRN,
    resolves spatial + Google Solar, and **persists** them (``epc_property`` /
    ``property_details_spatial`` / ``solar``) — then Baseline, then Modelling.
    ``run_modelling_e2e`` does ingestion inline and only models.
  * **There is no inspect-only mode**: the stages persist as they go (ADR-0012),
    so any run writes to the DB. This script is gated behind ``--confirm``; without
    it the script previews what it would do and exits.
  * **The modelling batch is all-or-nothing**: each stage commits once per batch,
    so one Property raising aborts the whole batch (no per-Property recovery like
    ``run_modelling_e2e``). Make sure the inputs are clean first.

Measure scoping comes **only from the Scenario's exclusions** — the pipeline
threads no ``--measures`` override (issue #1130). So if the live ``material``
catalogue cannot price/represent a measure a Property is eligible for (today:
``secondary_heating_removal``, absent from the ``material.type`` enum), that
Property's modelling raises and aborts the batch. Exclude it on the Scenario
first, e.g.::

    UPDATE scenario SET exclusions = '{secondary_heating_removal}' WHERE id = 1266;

EPC Prediction (ADR-0031) is left **off** — its Landlord-Override attributes
reader is not wired here, so an EPC-less Property is not gap-filled.

Config + secrets are loaded exactly as ``run_modelling_e2e`` does: ``backend/.env``
for the DB creds (``DB_*``), the EPC Bearer token (``OPEN_EPC_API_TOKEN``), the
Google Solar key (``GOOGLE_SOLAR_API_KEY``) and the S3 bucket (``DATA_BUCKET``);
AWS creds from the ambient ``~/.aws`` profile. Run from the worktree root::

    # preview only (no writes): print what would run, then exit
    python -m scripts.run_first_run_e2e --scenario-ids 1266 --portfolio-id 785 \
        709634 709635 709636
    # actually run the full pipeline and persist (Ingestion -> Baseline -> Modelling)
    python -m scripts.run_first_run_e2e --scenario-ids 1266 --portfolio-id 785 \
        --confirm 709634 709635 709636
"""

from __future__ import annotations

import argparse
import os
import sys
from pathlib import Path
from uuid import uuid4

_REPO_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(_REPO_ROOT))  # worktree root first — avoid the import trap

from applications.ara_first_run.ara_first_run_trigger_body import (  # noqa: E402
    AraFirstRunTriggerBody,
)
from applications.ara_first_run.handler import (  # noqa: E402
    build_first_run_pipeline,
    dispatch_first_run,
)
from infrastructure.epc_client.epc_client_service import EpcClientService  # noqa: E402
from infrastructure.solar.google_solar_api_client import (  # noqa: E402
    GoogleSolarApiClient,
)
from repositories.geospatial.geospatial_s3_repository import (  # noqa: E402
    GeospatialS3Repository,
)
from repositories.postgres_unit_of_work import PostgresUnitOfWork  # noqa: E402
from scripts.e2e_common import (  # noqa: E402
    ENV_PATH,
    build_engine,
    load_env,
    s3_parquet_reader,
)
from sqlmodel import Session  # noqa: E402


def _parse_ids(raw: str) -> list[int]:
    """Parse a comma-separated id list (e.g. ``--scenario-ids 1266,1270``)."""
    return [int(token.strip()) for token in raw.split(",") if token.strip()]


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "property_ids", type=int, nargs="+", help="Property ids to run"
    )
    parser.add_argument(
        "--scenario-ids",
        required=True,
        help="comma-separated Scenario ids to model against (exclusions come "
        "from each Scenario)",
    )
    parser.add_argument(
        "--portfolio-id", type=int, required=True, help="portfolio id for the run"
    )
    parser.add_argument(
        "--confirm",
        action="store_true",
        default=False,
        help="actually run the pipeline and WRITE to the DB (default: preview only)",
    )
    args = parser.parse_args()

    scenario_ids = _parse_ids(args.scenario_ids)

    load_env(ENV_PATH)
    engine = build_engine()

    body = AraFirstRunTriggerBody(
        # task/sub_task drive the Lambda SubTask lifecycle only; running the
        # pipeline directly bypasses the @subtask_handler decorator, so synthetic
        # ids satisfy validation without touching the task tables.
        task_id=uuid4(),
        sub_task_id=uuid4(),
        portfolio_id=args.portfolio_id,
        property_ids=args.property_ids,
        scenario_ids=scenario_ids,
    )

    print(
        f"full AraFirstRunPipeline (Ingestion -> Baseline -> Modelling) · "
        f"{len(args.property_ids)} propertie(s) · scenarios {scenario_ids} · "
        f"portfolio {args.portfolio_id}"
    )
    if not args.confirm:
        print(
            "\nPREVIEW ONLY — no writes. This run WOULD fetch + persist EPC/"
            "spatial/solar, rebaseline, and model+persist Plans for:\n"
            f"  properties: {args.property_ids}\n"
            "Re-run with --confirm to execute. NOTE: the modelling batch is "
            "all-or-nothing; ensure each Scenario excludes any measure the live "
            "catalogue cannot price (e.g. secondary_heating_removal)."
        )
        return

    epc_fetcher = EpcClientService(os.environ["OPEN_EPC_API_TOKEN"])
    geospatial_repo = GeospatialS3Repository(
        s3_parquet_reader(os.environ["DATA_BUCKET"])
    )
    solar_fetcher = GoogleSolarApiClient(os.environ["GOOGLE_SOLAR_API_KEY"])

    pipeline = build_first_run_pipeline(
        unit_of_work=lambda: PostgresUnitOfWork(lambda: Session(engine)),
        epc_fetcher=epc_fetcher,
        geospatial_repo=geospatial_repo,
        solar_fetcher=solar_fetcher,
    )

    print("running... (Ingestion -> Baseline -> Modelling, persisting per stage)\n")
    dispatch_first_run(body.model_dump(), pipeline=pipeline)
    print("done — EPC/spatial/solar + Baseline + Plans persisted for the batch.")


if __name__ == "__main__":
    main()