from collections.abc import Iterable, Sequence

import sqlalchemy as sa
from pydantic import TypeAdapter

from dmp.adapters.db import Repository
from dmp.tiktok.dtos import TikTokUserAdAccountLabel
from dmp.tiktok.models import TikTokAdAccount

TikTokUserAdAccountLabelAssociationListMapper = TypeAdapter(
    list[TikTokUserAdAccountLabel]
)


class TikTokAdAccountRepository(Repository[TikTokAdAccount]):
    def find_by_external_ids(
        self, external_ids: Iterable[str]
    ) -> Sequence[TikTokAdAccount]:
        query = sa.select(TikTokAdAccount).where(
            TikTokAdAccount.external_id.in_(external_ids)
        )
        result = self.db.session.execute(query)
        return result.scalars().all()

    def exists_by_id(self, id_: str) -> bool:
        query = sa.select(sa.select(1).exists().where(TikTokAdAccount.id == id_))
        result = self.db.session.execute(query)
        return bool(result.scalar_one())

    def count_by_ids(self, ids: list[str]) -> int:
        query = sa.select(sa.func.count()).where(TikTokAdAccount.id.in_(ids))
        result = self.db.session.execute(query)
        return result.scalar_one()

    def find_by_identity_id_and_user_id_for_assigment(
        self, identity_id: str, user_id: str
    ) -> list[TikTokUserAdAccountLabel]:
        query = self.db.query_from_template(
            "tiktok-ad-account/find-by-identity-id-and-user-id-for-assigment.sql",
            context={
                "identity_id": identity_id,
                "user_id": user_id,
            },
        )
        result = self.db.session.execute(query)
        return TikTokUserAdAccountLabelAssociationListMapper.validate_python(
            result.mappings().all()
        )