from __future__ import annotations

from typing import Optional

from sqlmodel import Session, col, delete, select

from domain.property_baseline.property_baseline_performance import PropertyBaselinePerformance
from infrastructure.postgres.property_baseline_performance_table import (
    PropertyBaselinePerformanceModel,
)
from repositories.property_baseline.property_baseline_repository import PropertyBaselineRepository


class PropertyBaselinePostgresRepository(PropertyBaselineRepository):
    """Maps PropertyBaselinePerformance to/from the ``property_baseline_performance`` table."""

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

    def save(self, baseline: PropertyBaselinePerformance, property_id: int) -> int:
        # Idempotent on property_id: a re-run (or re-score) replaces the row
        # rather than hitting the unique constraint (ADR-0012).
        self._session.exec(  # type: ignore[call-overload]
            delete(PropertyBaselinePerformanceModel).where(
                col(PropertyBaselinePerformanceModel.property_id) == property_id
            )
        )
        row = PropertyBaselinePerformanceModel.from_domain(baseline, property_id)
        self._session.add(row)
        self._session.flush()
        if row.id is None:
            raise ValueError("property_baseline_performance row did not receive an id")
        return row.id

    def get_for_property(
        self, property_id: int
    ) -> Optional[PropertyBaselinePerformance]:
        row = self._session.exec(
            select(PropertyBaselinePerformanceModel).where(
                PropertyBaselinePerformanceModel.property_id == property_id
            )
        ).first()
        return row.to_domain() if row is not None else None