from __future__ import annotations

from typing import Any, Optional

from sqlmodel import Session, col, select

from infrastructure.postgres.solar_table import SolarRow
from repositories.solar.solar_repository import SolarRepository


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

    def save(
        self, uprn: int, *, longitude: float, latitude: float, insights: dict[str, Any]
    ) -> None:
        existing = self._session.exec(
            select(SolarRow).where(SolarRow.uprn == uprn)
        ).first()
        if existing is None:
            self._session.add(
                SolarRow(
                    uprn=uprn,
                    longitude=longitude,
                    latitude=latitude,
                    google_api_response=insights,
                )
            )
        else:
            existing.longitude = longitude
            existing.latitude = latitude
            existing.google_api_response = insights
            self._session.add(existing)

    def get(self, uprn: int) -> Optional[dict[str, Any]]:
        row = self._session.exec(
            select(SolarRow).where(SolarRow.uprn == uprn)
        ).first()
        return row.google_api_response if row is not None else None

    def get_many(self, uprns: list[int]) -> dict[int, Optional[dict[str, Any]]]:
        """Stored insights for many UPRNs in one query — the batch form of
        ``get``. The returned dict has an entry for every requested UPRN; a UPRN
        with no stored row maps to None (exactly as ``get`` returns)."""
        rows = self._session.exec(
            select(SolarRow).where(col(SolarRow.uprn).in_(uprns))
        ).all()
        stored: dict[int, Optional[dict[str, Any]]] = {
            row.uprn: row.google_api_response for row in rows
        }
        return {uprn: stored.get(uprn) for uprn in uprns}