from __future__ import annotations

from abc import ABC, abstractmethod
from types import TracebackType
from typing import Optional

from repositories.property_baseline.property_baseline_repository import PropertyBaselineRepository
from repositories.epc.epc_repository import EpcRepository
from repositories.plan.plan_repository import PlanRepository
from repositories.product.product_repository import ProductRepository
from repositories.property.property_repository import PropertyRepository
from repositories.scenario.scenario_repository import ScenarioRepository
from repositories.solar.solar_repository import SolarRepository
from repositories.spatial.spatial_repository import SpatialRepository


class UnitOfWork(ABC):
    """A single batch transaction across the DB-backed repos (ADR-0012).

    A context manager that exposes the repos bound to one session. A stage runs
    its whole batch inside one unit and calls ``commit()`` once; leaving the
    block without committing — including via an exception — rolls back, so a
    failed batch persists nothing and the subtask fails noisily.

    The non-DB dependencies (EPC/Solar fetchers, the geospatial S3 repo, the
    Rebaseliner) are *not* part of the unit — only transactional DB work is.
    """

    property: PropertyRepository
    epc: EpcRepository
    solar: SolarRepository
    # Per-UPRN cache of the OS spatial reference data, written in Ingestion
    # alongside the EPC/solar enrichments (ADR-0020).
    spatial: SpatialRepository
    property_baseline: PropertyBaselineRepository
    # Modelling-stage repos (ADR-0017): read the Scenario, read the Product
    # catalogue, write the Plan + its Plan Measures — all on the one session.
    scenario: ScenarioRepository
    product: ProductRepository
    plan: PlanRepository

    @abstractmethod
    def commit(self) -> None: ...

    @abstractmethod
    def rollback(self) -> None: ...

    def __enter__(self) -> "UnitOfWork":
        return self

    def __exit__(
        self,
        exc_type: Optional[type[BaseException]],
        exc: Optional[BaseException],
        tb: Optional[TracebackType],
    ) -> None:
        # Roll back whatever was not explicitly committed (a no-op after a
        # successful commit). All-or-nothing per batch.
        self.rollback()