from collections.abc import Sequence

import sqlalchemy as sa
from sqlalchemy.orm import selectinload

from dmp.adapters.db import Repository
from dmp.google.models import GoogleUserConnection, GoogleUserConnectionAdAccount


class GoogleUserConnectionRepository(Repository[GoogleUserConnection]):
    default_options = [
        selectinload(GoogleUserConnection.connection_ad_accounts).joinedload(
            GoogleUserConnectionAdAccount.ad_account
        )
    ]

    def find_by_identity_id(self, identity_id: str) -> Sequence[GoogleUserConnection]:
        query = (
            sa.select(GoogleUserConnection)
            .where(GoogleUserConnection.identity_id == identity_id)
            .options(*self.default_options)
        )
        result = self.db.session.execute(query)
        return result.scalars().all()

    def find_by_identity_ids(
        self, identity_ids: Sequence[str]
    ) -> Sequence[GoogleUserConnection]:
        query = (
            sa.select(GoogleUserConnection)
            .where(GoogleUserConnection.identity_id.in_(identity_ids))
            .options(*self.default_options)
        )
        result = self.db.session.execute(query)
        return result.scalars().all()

    def get_by_identity_id_and_user_id(
        self, identity_id: str, user_id: str
    ) -> GoogleUserConnection | None:
        query = (
            sa.select(GoogleUserConnection)
            .where(
                GoogleUserConnection.identity_id == identity_id,
                GoogleUserConnection.user_id == user_id,
            )
            .options(*self.default_options)
        )
        result = self.db.session.execute(query)
        return result.scalar_one_or_none()

    def find_active_by_audience_and_ad_account_id(
        self, *, audience_id: str, ad_account_id: str, identity_id: str
    ) -> Sequence[GoogleUserConnection]:
        query = self.db.query_from_template(
            "google-user-connection/find-active-by-audience-id-and-ad-account-id.sql",
            context={
                "audience_id": audience_id,
                "ad_account_id": ad_account_id,
                "identity_id": identity_id,
            },
        )
        result = self.db.session.execute(
            sa.select(GoogleUserConnection)
            .from_statement(query)
            .options(*self.default_options)
        )
        return result.scalars().all()