Model/repositories/plan/plan_postgres_repository.py

from __future__ import annotations

from typing import Any

from sqlalchemy import insert as _sa_insert
from sqlmodel import Session, col, select, update

from domain.modelling.plan import Plan
from infrastructure.postgres.modelling import PlanModel, RecommendationModel
from repositories.plan.plan_repository import PlanRepository, PlanSaveRequest


def _col_values(model: Any, exclude: frozenset[str] = frozenset()) -> dict[str, Any]:
    """Extract column-keyed values from a SQLModel instance for Core INSERT."""
    return {
        c.name: getattr(model, c.name)
        for c in model.__table__.c
        if c.name not in exclude
    }


class PlanPostgresRepository(PlanRepository):
    """Maps a Plan and its Plan Measures onto the live ``plan`` /
    ``recommendation`` tables (ADR-0017). Does not commit — the Unit of Work
    owns the transaction (ADR-0012).

    A re-run INSERTs a fresh Plan rather than deleting the prior one (the cascade
    delete was slow); when the new Plan is the default it demotes any prior
    default Plan for the same property — across all Scenarios — to
    ``is_default=False``, so readers can select the property's one current Plan
    via ``is_default=True``."""

    def __init__(self, session: Session) -> None:
        self._session = session

    def save(
        self,
        plan: Plan,
        *,
        property_id: int,
        scenario_id: int,
        portfolio_id: int,
        is_default: bool,
    ) -> int:
        return self.save_batch(
            [PlanSaveRequest(plan, property_id=property_id, scenario_id=scenario_id, portfolio_id=portfolio_id, is_default=is_default)]
        )[0]

    def property_ids_with_default_plans(self, property_ids: list[int]) -> set[int]:
        if not property_ids:
            return set()
        rows = self._session.exec(
            select(col(PlanModel.property_id))
            .distinct()
            .where(
                col(PlanModel.property_id).in_(property_ids),
                col(PlanModel.is_default).is_(True),
            )
        ).all()
        return {int(row) for row in rows}

    def save_batch(self, requests: list[PlanSaveRequest]) -> list[int]:
        """Persist all Plans in three statements regardless of batch size.

        1. One demote UPDATE (only when any request has ``is_default=True``).
        2. One bulk plan INSERT with RETURNING to capture ids positionally.
        3. One bulk recommendation INSERT (skipped when no measures exist).
        """
        if not requests:
            return []

        # Demote prior default Plans for every property in the batch that is
        # receiving a new default Plan — one UPDATE for the whole batch,
        # across ALL scenarios: readers select a property's default by
        # portfolio + is_default alone, so the invariant is one default per
        # property. Scenario-scoped demotion left a promoted first Plan
        # (non-default scenario) defaulted forever alongside the default
        # scenario's later Plan (#1490).
        default_pids = [r.property_id for r in requests if r.is_default]
        if default_pids:
            self._session.exec(  # type: ignore[call-overload]
                update(PlanModel)
                .where(col(PlanModel.property_id).in_(default_pids))
                .values(is_default=False)
            )

        # Bulk INSERT all plan rows; capture returned ids positionally.
        plan_rows = [
            _col_values(
                PlanModel.from_domain(
                    r.plan,
                    property_id=r.property_id,
                    scenario_id=r.scenario_id,
                    portfolio_id=r.portfolio_id,
                    is_default=r.is_default,
                ),
                exclude=frozenset({"id"}),
            )
            for r in requests
        ]
        returned = self._session.execute(  # type: ignore[deprecated]
            _sa_insert(PlanModel).returning(PlanModel.__table__.c["id"]),  # type: ignore[attr-defined]
            plan_rows,
        ).all()
        plan_ids = [row[0] for row in returned]

        # Accumulate recommendation rows across all requests; properties with
        # zero measures contribute nothing (no special-casing needed).
        rec_rows = [
            _col_values(
                RecommendationModel.from_domain(
                    measure, property_id=r.property_id, plan_id=plan_id
                ),
                exclude=frozenset({"id"}),
            )
            for r, plan_id in zip(requests, plan_ids)
            for measure in r.plan.measures
        ]
        if rec_rows:
            self._session.execute(  # type: ignore[deprecated]
                _sa_insert(RecommendationModel), rec_rows
            )

        return plan_ids