Model/backend/app/modelling/router.py

"""The Modelling Run Distributor: POST /v1/modelling/trigger-run (ADR-0055).

Accepts a portfolio-scoped modelling request expressed as filters, resolves
them to a concrete property set (ADR-0056), pre-creates one batch sub_task per
SQS message under the app-owned task, and fans the batches out to the
modelling_e2e workers. Never models synchronously; owns nothing after the
fan-out — progress and terminal state roll up from the workers.
"""

import json
from collections.abc import Callable, Iterator
from typing import Any, cast
from uuid import uuid4

import boto3
from fastapi import APIRouter, Depends, HTTPException
from sqlalchemy import text
from sqlmodel import Session

from backend.app.config import get_settings
from backend.app.db.connection import db_engine
from backend.app.dependencies import validate_token
from backend.app.modelling.batching import pack_postcode_batches
from backend.app.modelling.property_filters import (
    FilteredProperty,
    resolve_filtered_property_ids,
)
from backend.app.modelling.run_tasks import ModellingRunTasks
from backend.app.modelling.schemas import TriggerRunRequest

# Sends pre-serialised message bodies to the modelling_e2e queue. A seam so
# tests record bodies instead of calling AWS.
MessageSender = Callable[[list[str]], None]


def get_session() -> Iterator[Session]:
    with Session(db_engine) as session:
        yield session


def get_message_sender() -> MessageSender:
    settings = get_settings()
    client: Any = cast(Any, boto3.client("sqs", settings.AWS_DEFAULT_REGION))  # pyright: ignore[reportUnknownMemberType]
    queue_url = settings.MODELLING_E2E_SQS_URL

    def send(bodies: list[str]) -> None:
        # send_message_batch caps at 10 entries per call — chunk accordingly.
        for start in range(0, len(bodies), 10):
            chunk = bodies[start : start + 10]
            client.send_message_batch(
                QueueUrl=queue_url,
                Entries=[
                    {"Id": str(index), "MessageBody": body}
                    for index, body in enumerate(chunk)
                ],
            )

    return send


router = APIRouter(
    prefix="/modelling",
    tags=["modelling"],
    dependencies=[Depends(validate_token)],
)


@router.post("/trigger-run", status_code=202)
async def trigger_run(
    body: TriggerRunRequest,
    session: Session = Depends(get_session),
    send_messages: MessageSender = Depends(get_message_sender),
) -> dict[str, str]:
    run_tasks = ModellingRunTasks(session)
    if run_tasks.already_distributed(body.task_id):
        raise HTTPException(
            status_code=409,
            detail=(
                f"Task {body.task_id} already has sub_tasks — it has been "
                "distributed. Check its progress instead of re-triggering."
            ),
        )

    scenario_rows = session.connection().execute(
        text("SELECT id, portfolio_id FROM scenario WHERE id = ANY(:ids)"),
        {"ids": body.scenario_ids},
    )
    portfolio_by_scenario = {int(row[0]): row[1] for row in scenario_rows}
    invalid = [
        scenario_id
        for scenario_id in body.scenario_ids
        if portfolio_by_scenario.get(scenario_id) != body.portfolio_id
    ]
    if invalid:
        raise HTTPException(
            status_code=400,
            detail=(
                f"Scenarios {invalid} do not belong to portfolio "
                f"{body.portfolio_id}."
            ),
        )

    properties: list[FilteredProperty] = resolve_filtered_property_ids(
        session, body.portfolio_id, body.filters
    )
    if not properties:
        # A task with zero sub_tasks could never roll up to complete; the
        # app's preview shows the same zero from the same rule (ADR-0056).
        raise HTTPException(
            status_code=400,
            detail=(
                f"The filters resolve to no properties in portfolio "
                f"{body.portfolio_id} — nothing to distribute."
            ),
        )
    batches = pack_postcode_batches(properties)

    # Pre-create one sub_task per (scenario, batch) message under the
    # app-owned task, each holding its exact message payload — the fixed
    # progress denominator and the batch's re-run recipe (ADR-0055).
    messages: list[dict[str, Any]] = []
    for scenario_id in body.scenario_ids:
        for batch in batches:
            messages.append(
                {
                    "task_id": str(body.task_id),
                    "subtask_id": str(uuid4()),
                    "property_ids": [p.property_id for p in batch],
                    "portfolio_id": body.portfolio_id,
                    "scenario_id": scenario_id,
                    # ADR-0055 pinned flags: live EPC fetch, live prediction,
                    # solar fetched only where no stored row exists.
                    "refetch_epc": True,
                    "repredict_epc": True,
                    "refetch_solar": True,
                    "dry_run": False,
                }
            )

    run_tasks.create_batch_subtasks(body.task_id, messages)

    send_messages([json.dumps(message) for message in messages])
    return {"message": "Modelling Run distributed"}