# pylint: disable=too-many-lines
import asyncio
import gzip
import json
import time
from datetime import datetime, timedelta, timezone
from functools import cache, partial
from itertools import groupby
from operator import attrgetter, itemgetter
from typing import Any, Optional

import httpx
import orjson
import sentry_sdk
import structlog
from datadog.dogstatsd.base import statsd
from sqlalchemy import inspect, select, update
from sqlalchemy.dialects.postgresql import insert
from sqlalchemy.engine.row import Row
from sqlalchemy.ext.asyncio import AsyncEngine
from sqlalchemy.sql.expression import and_, or_

from dapd_api_scraper.aws import AWSCredentials, S3AsyncClient
from dapd_api_scraper.entities import (
    AppleMusicAlbumEntity,
    AppleMusicArtistEntity,
    AppleMusicPlaylistEntity,
    AppleMusicTrackEntity,
    SpotifyAlbumEntity,
    SpotifyArtistEntity,
    SpotifyPlaylistEntity,
    SpotifyTrackEntity,
)
from dapd_api_scraper.etl_models import (
    AppleMusicAlbum,
    AppleMusicArtist,
    AppleMusicPlaylist,
    AppleMusicTrack,
    PlaylistFollowers,
    SpotifyAlbum,
    SpotifyArtist,
    SpotifyPlaylist,
    SpotifyTrack,
)
from dapd_api_scraper.etldb_models import (
    DimAlbum,
    DimAlbumArtist,
    DimArtist,
    DimDsp,
    DimMarket,
    DimPlaylist,
    DimPlaylistMeta,
    DimTrack,
    DimTrackAlbum,
    DimTrackArtist,
    FactArtistFollowers,
    FactPlaylistFollowers,
    PlaylistOwner,
    PlaylistOwnerCategory,
)
from dapd_api_scraper.utils import (
    ShutdownHandler,
    async_cache,
    async_execute,
    async_retry,
    chunks,
    deep_chunks,
)
from dapd_api_scraper.workflowdb_models import (
    Album,
    Artist,
    BuzzCategory,
    BuzzUser,
    Playlist,
    Track,
)


EtlModelType = (
    SpotifyPlaylist
    | SpotifyAlbum
    | SpotifyArtist
    | SpotifyTrack
    | AppleMusicAlbum
    | AppleMusicArtist
    | AppleMusicPlaylist
    | AppleMusicTrack
)

EntityType = (
    AppleMusicAlbumEntity
    | AppleMusicArtistEntity
    | AppleMusicPlaylistEntity
    | AppleMusicTrackEntity
    | SpotifyAlbumEntity
    | SpotifyArtistEntity
    | SpotifyPlaylistEntity
    | SpotifyTrackEntity
)

WorkflowDbModelType = Album | Artist | BuzzCategory | BuzzUser | Playlist | Track


async def process_entity(
    logger: structlog.BoundLogger,
    shutdown: ShutdownHandler,
    workflowdb_engine: AsyncEngine,
    etldb_engine: AsyncEngine,
    entity: EntityType,
) -> None:
    """Processes given entity until SIGTERM signal recieved or
    working time exceed 1 hour"""
    logger = logger.bind(data_source=entity.data_source, entity=entity.entity_type)

    s3_client = S3AsyncClient(
        entity.config.aws_region,
        AWSCredentials(entity.config.aws_credentials_url),
        entity.config.raw_bucket,
    )

    start_time = time.time()

    process_time = 0.0

    while (not shutdown.signal_recieved) and (process_time < 3600):
        records = await get_records(workflowdb_engine, entity)
        grouped_records = group_records_by_storefront(records)

        if not grouped_records:
            logger.info(
                "No expired records, repeating after 5 second",
            )
            await asyncio.sleep(5)

        semaphore = asyncio.Semaphore(entity.threads_count)

        results = await asyncio.gather(
            *map(
                partial(
                    process_records,
                    logger,
                    shutdown,
                    semaphore,
                    workflowdb_engine,
                    etldb_engine,
                    s3_client,
                    entity,
                ),
                deep_chunks(grouped_records, entity.chunk_size),
            ),
            return_exceptions=True,
        )

        for result in results:
            if isinstance(result, BaseException):
                sentry_sdk.capture_exception(result)
                logger.error("Unhandled error", exc_info=result)

        process_time = time.time() - start_time

        # healthcheck marks
        with open("healthcheck.json", "r", encoding="utf-8") as file_object:
            healthcheck = json.load(file_object)

        healthcheck[f"{entity.data_source}_{entity.entity_type}"] = int(time.time())
        with open("healthcheck.json", "w", encoding="utf-8") as file_object:
            json.dump(healthcheck, file_object, indent=2)


def group_records_by_storefront(records: list[Row]) -> list[list[Row]]:
    """Groups records by storefront to try to process them by batches"""
    records.sort(key=attrgetter("storefront"))
    grouped_records = [list(v) for _, v in groupby(records, attrgetter("storefront"))]
    return grouped_records


async def process_records(  # pylint: disable=too-many-arguments
    logger: structlog.BoundLogger,
    shutdown: ShutdownHandler,
    semaphore: asyncio.Semaphore,
    workflowdb_engine: AsyncEngine,
    etldb_engine: AsyncEngine,
    s3_client: S3AsyncClient,
    entity: EntityType,
    records: list[Row],
) -> None:
    """Processes given records"""

    async with semaphore:
        logger = logger.bind(
            records=[record.id for record in records],
            storefront=records[0].storefront,
        )

        if shutdown.signal_recieved:
            return

        try:
            response = await get_api_response(entity, records)
            data = orjson.loads(response.content)

            if entity.data_source == "spotify" and entity.entity_type == "playlist":
                if records[0].save_tracklist:
                    data = await extend_spotify_playlist_with_paginated_tracks(logger, entity, data)
                else:
                    data.pop("tracks", None)
            elif entity.data_source == "apple_music" and entity.entity_type == "playlist":
                data = await extend_apple_music_playlist_with_paginated_tracks(entity, data)

            results = await asyncio.gather(
                store_data_to_etldb(logger, etldb_engine, entity, records, data),
                store_data_to_workflowdb(workflowdb_engine, entity, records, data=data),
                store_data_to_s3(s3_client, entity, records, data),
                return_exceptions=True,
            )

            for result in results:
                if isinstance(result, BaseException):
                    raise result

            logger.info(f"Processed {len(records)} records")

        except httpx.HTTPStatusError as exc:
            statsd.increment(
                "dapd-api-scraper.api_errors",
                tags=[
                    f"dsp:{entity.data_source}",
                    f"entity:{entity.entity_type}",
                    f"status_code:{exc.response.status_code}",
                ],
            )
            logger.warning(
                f"HTTP Error {exc.response.status_code}",
            )
            if exc.response.status_code == 429:
                await asyncio.sleep(5)
            if exc.response.status_code in (400, 404):
                await store_data_to_workflowdb(
                    workflowdb_engine,
                    entity,
                    records,
                    not_found=True,
                )
            if (
                exc.response.status_code == 502
                and entity.data_source == "spotify"
                and entity.entity_type == "playlist"
            ):
                await store_data_to_workflowdb(
                    workflowdb_engine,
                    entity,
                    records,
                    not_found=True,
                )
        except httpx.RemoteProtocolError:
            logger.warning("Remote protocol error")
        except httpx.ConnectTimeout:
            logger.warning("Connect timeout")
        except httpx.ReadTimeout:
            logger.warning("Read timeout")
        except httpx.ConnectError:
            logger.warning("Connect error")
        except httpx.ReadError:
            logger.warning("Read error")
        except Exception as exc:
            sentry_sdk.capture_exception(exc)
            logger.error("Unhandled error", exc_info=exc)
        finally:
            statsd.increment(
                "dapd-api-scraper.records_processed",
                value=len(records),
                tags=[f"dsp:{entity.data_source}", f"entity:{entity.entity_type}"],
            )


async def extend_spotify_playlist_with_paginated_tracks(
    logger: structlog.BoundLogger,
    entity: EntityType,
    data: dict[str, Any],
) -> Any:
    """Return extended playlists with tracks in case if tracks count
    in playlist more than 100, otherwise unchanged playlist will be returned
    """
    total = data.get("tracks", {}).get("total", 0)

    if not total or total <= 100:
        return data

    tracks_chunks = await asyncio.gather(
        *map(partial(get_tracks, logger, entity, data), range(100, total, 100)),
        return_exceptions=True,
    )

    for tracks_chunk in tracks_chunks:
        if isinstance(tracks_chunk, BaseException):
            raise tracks_chunk

    for tracks in sorted(tracks_chunks, key=itemgetter(1)):
        data["tracks"]["items"].extend(tracks[0])

    return data


async def extend_apple_music_playlist_with_paginated_tracks(
    entity: EntityType,
    data: dict[str, Any],
) -> dict[str, Any]:
    """Return extended playlists with tracks in case if tracks count
    in playlist more than 100, otherwise unchanged playlist will be returned
    """

    for playlist in data["data"]:
        additional_tracks_list = []
        tracks = playlist["relationships"]["tracks"]
        while tracks.get("next"):
            path = tracks.get("next").replace("/v1/catalog", "")
            response = await get_additional_apple_music_tracks(entity, path)
            response.raise_for_status()
            tracks = orjson.loads(response.content)
            additional_tracks_list.extend(tracks["data"])

        playlist["relationships"]["tracks"]["data"].extend(additional_tracks_list)

    return data


@async_retry(delay=1, max_retries=3)
async def get_additional_apple_music_tracks(entity: EntityType, path: str) -> httpx.Response:
    api_client = get_api_client(entity)
    auth_headers = entity.get_auth_headers()
    with statsd.timed(
        "dapd-api-scraper.api_response_time",
        tags=[
            f"dsp:{entity.data_source}",
            "entity:playlist_tracks",
        ],
        use_ms=True,
    ):
        tracks = await api_client.get(path, headers=auth_headers)
    return tracks


@async_retry(delay=1, max_retries=3)
async def get_tracks(
    logger: structlog.BoundLogger,
    entity: EntityType,
    data: dict[str, Any],
    offset: int,
) -> tuple[list[dict[str, Any]], int]:
    """Returns spotify playlist's tracks with the given offset"""
    try:
        api_client = get_api_client(entity)
        path = get_path_for_additional_spotify_tracks(data, offset)
        auth_headers = entity.get_auth_headers()
        with statsd.timed(
            "dapd-api-scraper.api_response_time",
            tags=[
                f"dsp:{entity.data_source}",
                "entity:playlist_tracks",
            ],
            use_ms=True,
        ):
            response = await api_client.get(path, headers=auth_headers)
        response.raise_for_status()
        tracks = orjson.loads(response.content)["items"]
        logger.info(f"Acquired additional {len(tracks)} tracks")

        return tracks, offset
    except httpx.HTTPStatusError as exc:
        statsd.increment(
            "dapd-api-scraper.api_errors",
            tags=[
                f"dsp:{entity.data_source}",
                "entity:playlist_tracks",
                f"status_code:{exc.response.status_code}",
            ],
        )
        logger.warning(
            f"HTTP Error {exc.response.status_code}",
        )
        raise
    except Exception as exc:
        sentry_sdk.capture_exception(exc)
        logger.error("Unhandled error", exc_info=exc)
        raise


def get_path_for_additional_spotify_tracks(
    data: dict[str, Any],
    offset: int,
) -> str:
    """Path constructor for additional tracks for playlist"""
    return f'playlists/{data["id"]}/tracks?offset={offset}&limit=100'


def get_records_query(model, data_source: str, limit: int) -> Any:
    """Returns expired records query from workflowdb which should be processed"""
    if hasattr(model, "buzz_user_username"):
        query = (
            select(
                model,
                model.buzz_user_username.label("bu_username"),
                BuzzUser.display_name.label("bu_display_name"),
                BuzzUser.storefront.label("bu_country_code"),
                BuzzCategory.id.label("bu_category_id"),
                BuzzCategory.name.label("bc_name"),
            )
            .join(BuzzUser, model.buzz_user_username == BuzzUser.username, isouter=True)
            .join(BuzzCategory, BuzzUser.buzz_category_id == BuzzCategory.id, isouter=True)
        )
    else:
        query = select(model)

    query = (
        query.where(model.data_source == data_source)
        .where(model.expired_at < datetime.now(timezone.utc))
        .where(
            or_(
                model.is_removed.is_not(True),
                and_(
                    model.is_removed.is_(True),
                    model.is_removed_expired_at < datetime.now(timezone.utc),
                ),
            )
        )
        .order_by(model.expired_at.asc())
        .limit(limit)
    )

    return query


async def get_records(workflowdb_engine: AsyncEngine, entity: EntityType) -> Any:
    """Returns expired records from workflowdb which should be processed"""
    query = get_records_query(
        entity.workflowdb_model,
        entity.data_source,
        entity.db_get_limit,
    )
    result = await async_execute(workflowdb_engine, query)
    return result.fetchall()


async def get_api_response(entity: EntityType, records: list[Row]) -> httpx.Response:
    """Returns performed response"""
    api_client = get_api_client(entity)
    path = entity.get_api_path_for_records(records)
    auth_headers = entity.get_auth_headers()
    with statsd.timed(
        "dapd-api-scraper.api_response_time",
        tags=[
            f"dsp:{entity.data_source}",
            f"entity:{entity.entity_type}",
        ],
        use_ms=True,
        sample_rate=1,
    ):
        response = await api_client.get(path, headers=auth_headers)
        statsd.increment(
            "dapd-api-scraper.api_request_count",
            tags=[
                f"dsp:{entity.data_source}",
                f"entity:{entity.entity_type}",
            ],
        )
    response.raise_for_status()
    return response


@cache
def get_api_client(entity: EntityType) -> httpx.AsyncClient:
    """Creates new httpx.Client instance with base_url and timeout
    specified in Entity argument
    """
    return httpx.AsyncClient(base_url=entity.api_url, timeout=entity.api_timeout, http2=False)


async def store_data_to_workflowdb(
    workflowdb_engine: AsyncEngine,
    entity: EntityType,
    records: list[Row],
    data: Optional[Any] = None,
    not_found: bool = False,
) -> None:
    """Stores processed records to workflowdb"""
    if not_found:
        await update_records_in_workflowdb(
            workflowdb_engine,
            entity,
            records,
            data,
            not_found=not_found,
        )
    else:
        results = await asyncio.gather(
            update_records_in_workflowdb(
                workflowdb_engine, entity, records, data, not_found=not_found
            ),
            store_tracks_to_workflowdb(workflowdb_engine, entity, records, data),
            store_albums_to_workflowdb(workflowdb_engine, entity, records, data),
            store_artists_to_workflowdb(workflowdb_engine, entity, records, data),
            return_exceptions=True,
        )

        for result in results:
            if isinstance(result, BaseException):
                raise result


async def update_records_in_workflowdb(
    workflowdb_engine: AsyncEngine,
    entity: EntityType,
    records: list[Row],
    data: Any,
    not_found=False,
):
    now = datetime.now(timezone.utc)
    returned_entities_ids = entity.get_entities_ids_from_response(data)
    model = entity.workflowdb_model
    for record in records:
        payload = {
            "expired_at": now + timedelta(minutes=record.ttl_minutes),
            "updated_at": now,
            "is_removed": False,
            "is_removed_expired_at": None,
        }
        if (not_found) or (not data) or (record.id not in returned_entities_ids):
            payload["is_removed"] = True
            payload["is_removed_expired_at"] = now + min(
                timedelta(minutes=record.ttl_minutes),
                timedelta(days=entity.is_removed_expired_ttl_days),
            )
            payload["expired_at"] = now + min(
                timedelta(minutes=record.ttl_minutes),
                timedelta(days=entity.is_removed_expired_ttl_days),
            )
        elif hasattr(record, "snapshot_id"):
            api_data_for_record = entity.get_record_from_api_response(record, data)
            if api_data_for_record:
                payload["snapshot_id"] = api_data_for_record.get("snapshot_id")
        query = (
            update(model)
            .values(payload)
            .where(
                and_(
                    model.id == record.id,
                    model.data_source == record.data_source,
                    model.storefront == record.storefront,
                )
            )
        )
        await async_execute(workflowdb_engine, query)


async def store_tracks_to_workflowdb(
    workflowdb_engine: AsyncEngine, entity: EntityType, records: list[Row], data: Any
) -> None:
    """
    Stores tracks if they exists for entity"""
    tracks = entity.get_tracks_from_data(records, data)

    if tracks:
        existing_tracks = []
        for tracks_chunk in chunks([track["id"] for track in tracks], 10000):
            existing_tracks_query = select(Track.id).where(
                and_(
                    Track.id.in_(tracks_chunk),
                    Track.storefront == records[0].storefront,
                    Track.data_source == records[0].data_source,
                )
            )
            result = await async_execute(workflowdb_engine, existing_tracks_query)
            existing_tracks.extend([track.id for track in result.fetchall()])

        if existing_tracks:
            tracks = [track for track in tracks if track["id"] not in existing_tracks]

        if not tracks:
            return

        for tracks_chunk in chunks(tracks, 1000):
            tracks_query = insert(Track).values(tracks_chunk).on_conflict_do_nothing()
            await async_execute(workflowdb_engine, tracks_query)


async def store_albums_to_workflowdb(
    workflowdb_engine: AsyncEngine, entity: EntityType, records: list[Row], data: Any
) -> None:
    """
    Stores albums if they exists for entity"""
    albums = entity.get_albums_from_data(records, data)
    if albums:
        existing_albums_query = select(Album.id).where(
            and_(
                Album.id.in_([album["id"] for album in albums]),
                Album.storefront == records[0].storefront,
                Album.data_source == records[0].data_source,
            )
        )

        result = await async_execute(workflowdb_engine, existing_albums_query)
        existing_albums_ids = [album.id for album in result.fetchall()]

        if existing_albums_ids:
            albums = [album for album in albums if album["id"] not in existing_albums_ids]

        if not albums:
            return

        for albums_chunk in chunks(albums, 1000):
            albums_query = insert(Album).values(albums_chunk).on_conflict_do_nothing()
            await async_execute(workflowdb_engine, albums_query)


async def store_artists_to_workflowdb(
    workflowdb_engine: AsyncEngine, entity: EntityType, records: list[Row], data: Any
) -> None:
    """
    Stores artists if they exists for entity"""
    artists = entity.get_artists_from_data(records, data)
    if artists:
        existing_artists_query = select(Artist.id).where(
            and_(
                Artist.id.in_([artist["id"] for artist in artists]),
                Artist.storefront == records[0].storefront,
                Artist.data_source == records[0].data_source,
            )
        )

        result = await async_execute(workflowdb_engine, existing_artists_query)
        existing_artists_ids = [artist.id for artist in result.fetchall()]

        if existing_artists_ids:
            artists = [artist for artist in artists if artist["id"] not in existing_artists_ids]

        if not artists:
            return

        for artists_chunk in chunks(artists, 1000):
            artists_query = insert(Artist).values(artists_chunk).on_conflict_do_nothing()

            await async_execute(workflowdb_engine, artists_query)


async def store_data_to_etldb(
    logger: structlog.BoundLogger,
    etldb_engine: AsyncEngine,
    entity: EntityType,
    records: list[Row],
    data: Any,
) -> None:
    """Stores processed records to etldb"""

    excluded_fields = {
        "tracks",
        "artists",
        "albums",
        "playlist_metadata",
        "artist_followers",
        "playlist_tracks_dynamics",
        "track_albums",
        "album_artists",
        "track_artists",
        "playlist_owner",
        "playlist_owner_category",
    }

    market_id = await get_market_id_by_storefront(etldb_engine, records[0].storefront)
    dsp_id = await get_dsp_id_by_data_source(etldb_engine, entity.data_source)

    for record in records:
        exception_happened = False
        api_data_for_record = entity.get_record_from_api_response(record, data)
        if not api_data_for_record:
            continue

        if (
            record.data_source == "spotify"
            and entity.entity_type == "playlist"
            and record.save_tracklist is True
        ):
            fact_playlist_followers = PlaylistFollowers(
                market_id=market_id, dsp_id=dsp_id, **api_data_for_record
            )
            await store_fact_playlist_followers(etldb_engine, **fact_playlist_followers.dict())
        if hasattr(record, "snapshot_id") and record.snapshot_id == api_data_for_record.get(
            "snapshot_id"
        ):
            continue
        model = entity.etl_model(
            market_id=market_id,
            dsp_id=dsp_id,
            **{
                **record,
                **api_data_for_record,
                "playlist_type": record.type if hasattr(record, "type") else None,
                "updated_at": datetime.now(timezone.utc),
            },
        )

        if hasattr(model, "playlist_owner_category") and model.playlist_owner_category:
            await store_playlist_owner_category(
                etldb_engine, entity, **model.playlist_owner_category.dict()
            )

        if hasattr(model, "playlist_owner") and model.playlist_owner:
            await store_playlist_owner(etldb_engine, entity, **model.playlist_owner.dict())

        dumped_model = model.dict(exclude=excluded_fields)
        query = (
            insert(entity.etldb_model)
            .values({**dumped_model, "created_at": dumped_model["updated_at"]})
            .on_conflict_do_update(
                index_elements=get_model_pk(entity.etldb_model),
                set_=dumped_model,
            )
        )
        await async_execute(etldb_engine, query)

        results = await asyncio.gather(
            store_tracks_to_etldb(
                etldb_engine,
                model,
                excluded_fields,
                entity,
                records[0].storefront,
            ),
            store_playlist_metadata_to_etldb(etldb_engine, model, entity),
            store_albums_to_etldb(
                etldb_engine, model, excluded_fields, entity, records[0].storefront
            ),
            store_artists_to_etldb(etldb_engine, model, entity, records[0].storefront),
            return_exceptions=True,
        )

        for result in results:
            if isinstance(result, BaseException):
                exception_happened = True
                sentry_sdk.capture_exception(result)
                logger.error("Unhandled error", exc_info=result)

        if not exception_happened:
            results = await asyncio.gather(
                store_artist_followers_to_etldb(etldb_engine, model, entity),
                store_track_albums_to_etldb(etldb_engine, model, entity),
                store_album_artists_to_etldb(etldb_engine, model, entity),
                store_track_artists_to_etldb(etldb_engine, model, entity),
                return_exceptions=True,
            )

            for result in results:
                if isinstance(result, BaseException):
                    exception_happened = True
                    sentry_sdk.capture_exception(result)
                    logger.error("Unhandled error", exc_info=result)


@async_cache(timeout=86400)
async def store_playlist_owner_category(etldb_engine, entity, **playlist_owner_category):
    now = datetime.now(timezone.utc)
    playlist_owner_category["created_at"] = now
    playlist_owner_category["updated_at"] = now
    query = insert(PlaylistOwnerCategory).values(playlist_owner_category).on_conflict_do_nothing()
    await async_execute(etldb_engine, query)


@async_cache(timeout=3600)
async def store_playlist_owner(etldb_engine, entity, **playlist_owner):
    if playlist_owner.get("market_id"):
        playlist_owner["market_id"] = await get_market_id_by_storefront(
            etldb_engine,
            playlist_owner["market_id"].lower(),
        )
        now = datetime.now(timezone.utc)
        playlist_owner["created_at"] = now
        playlist_owner["updated_at"] = now

    query = (
        insert(PlaylistOwner)
        .values(**playlist_owner)
        .on_conflict_do_update(index_elements=get_model_pk(PlaylistOwner), set_=playlist_owner)
    )
    await async_execute(etldb_engine, query)


async def store_tracks_to_etldb(
    etldb_engine: AsyncEngine,
    model: EtlModelType,
    excluded_fields: set[str],
    entity: EntityType,
    storefront: str,
) -> None:
    if isinstance(model, (SpotifyPlaylist, AppleMusicPlaylist)) and model.tracks:
        tracks_to_write = model.dict(include={"tracks"})["tracks"]
        for track in tracks_to_write:
            for excluded_field in excluded_fields:
                track.pop(excluded_field, None)

        existing_tracks_ids = []

        for tracks_chunk in chunks(tracks_to_write, 1000):
            existing_tracks_query = select(DimTrack.track_id).where(
                and_(
                    DimTrack.track_id.in_([track["track_id"] for track in tracks_chunk]),
                    DimTrack.market_id
                    == await get_market_id_by_storefront(etldb_engine, storefront),
                    DimTrack.dsp_id
                    == await get_dsp_id_by_data_source(etldb_engine, entity.data_source),
                )
            )
            result = await async_execute(etldb_engine, existing_tracks_query)
            existing_tracks_ids.extend([track.track_id for track in result.fetchall()])

        if existing_tracks_ids:
            tracks_to_write = [
                track for track in tracks_to_write if track["track_id"] not in existing_tracks_ids
            ]

        if not tracks_to_write:
            return

        for tracks_chunk in chunks(tracks_to_write, 1000):
            for track in tracks_chunk:
                now = datetime.now(timezone.utc)
                track["created_at"] = now
                track["updated_at"] = now

            query = insert(DimTrack).values(tracks_chunk).on_conflict_do_nothing()
            await async_execute(etldb_engine, query)


async def store_playlist_metadata_to_etldb(
    etldb_engine: AsyncEngine,
    model: EtlModelType,
    entity: EntityType,
) -> None:
    if (
        isinstance(model, (SpotifyPlaylist, AppleMusicPlaylist))
        and model.playlist_metadata is not None
    ):
        query = insert(DimPlaylistMeta).values(model.playlist_metadata.dict())
        await async_execute(etldb_engine, query)


async def store_artist_followers_to_etldb(
    etldb_engine: AsyncEngine,
    model: EtlModelType,
    entity: EntityType,
) -> None:
    if isinstance(model, SpotifyArtist) and model.artist_followers:
        query = insert(FactArtistFollowers).values(model.artist_followers.dict())
        await async_execute(etldb_engine, query)


async def store_albums_to_etldb(
    etldb_engine: AsyncEngine,
    model: EtlModelType,
    excluded_fields: set[str],
    entity: EntityType,
    storefront: str,
) -> None:
    if isinstance(model, (SpotifyTrack, AppleMusicTrack)) and model.albums:
        albums_to_write = model.dict(include={"albums"})["albums"]
        for album in albums_to_write:
            for excluded_field in excluded_fields:
                album.pop(excluded_field, None)

        existing_albums_ids = []

        for albums_chunk in chunks(albums_to_write, 1000):
            existing_albums_query = select(DimAlbum.album_id).where(
                and_(
                    DimAlbum.album_id.in_([album["album_id"] for album in albums_chunk]),
                    DimAlbum.market_id
                    == await get_market_id_by_storefront(etldb_engine, storefront),
                    DimAlbum.dsp_id
                    == await get_dsp_id_by_data_source(etldb_engine, entity.data_source),
                )
            )
            result = await async_execute(etldb_engine, existing_albums_query)
            existing_albums_ids.extend([album.album_id for album in result.fetchall()])

        if existing_albums_ids:
            albums_to_write = [
                album for album in albums_to_write if album["album_id"] not in existing_albums_ids
            ]

        if not albums_to_write:
            return

        for albums_chunk in chunks(albums_to_write, 1000):
            for album in albums_chunk:
                now = datetime.now(timezone.utc)
                album["created_at"] = now
                album["updated_at"] = now

            query = insert(DimAlbum).values(albums_chunk).on_conflict_do_nothing()
            await async_execute(etldb_engine, query)


async def store_artists_to_etldb(
    etldb_engine: AsyncEngine,
    model: EtlModelType,
    entity: EntityType,
    storefront: str,
) -> None:
    if (
        isinstance(model, (SpotifyAlbum, SpotifyTrack, AppleMusicAlbum, AppleMusicTrack))
        and model.artists
    ):
        artists_to_write = model.dict(include={"artists"})["artists"]
        for artist in artists_to_write:
            artist.pop("artist_followers", None)

        existing_artists_ids = []

        for artists_chunk in chunks(artists_to_write, 1000):
            existing_artists_query = select(DimArtist.artist_id).where(
                and_(
                    DimArtist.artist_id.in_([artist["artist_id"] for artist in artists_to_write]),
                    DimArtist.market_id
                    == await get_market_id_by_storefront(etldb_engine, storefront),
                    DimArtist.dsp_id
                    == await get_dsp_id_by_data_source(etldb_engine, entity.data_source),
                )
            )

            result = await async_execute(etldb_engine, existing_artists_query)
            existing_artists_ids.extend([artist.artist_id for artist in result.fetchall()])

        if existing_artists_ids:
            artists_to_write = [
                artist
                for artist in artists_to_write
                if artist["artist_id"] not in existing_artists_ids
            ]

        if not artists_to_write:
            return

        for artists_chunk in chunks(artists_to_write, 1000):
            for artist in artists_chunk:
                now = datetime.now(timezone.utc)
                artist["created_at"] = now
                artist["updated_at"] = now

            query = insert(DimArtist).values(artists_chunk).on_conflict_do_nothing()

            await async_execute(etldb_engine, query)


async def store_track_albums_to_etldb(
    etldb_engine: AsyncEngine,
    model: EtlModelType,
    entity: EntityType,
) -> None:
    if isinstance(model, (SpotifyTrack, AppleMusicTrack)) and model.track_albums:
        track_albums_to_write = model.dict(include={"track_albums"})["track_albums"]

        for track_album in track_albums_to_write:
            now = datetime.now(timezone.utc)
            track_album["created_at"] = now
            track_album["updated_at"] = now

        query = (
            insert(DimTrackAlbum)
            .values(track_albums_to_write)
            .on_conflict_do_nothing(
                index_elements=get_model_pk(DimTrackAlbum),
            )
        )

        await async_execute(etldb_engine, query)


async def store_album_artists_to_etldb(
    etldb_engine: AsyncEngine,
    model: EtlModelType,
    entity: EntityType,
) -> None:
    if isinstance(model, (SpotifyAlbum, AppleMusicAlbum)) and model.album_artists:
        album_artists_to_write = model.dict(include={"album_artists"})["album_artists"]

        for album_artist in album_artists_to_write:
            now = datetime.now(timezone.utc)
            album_artist["created_at"] = now
            album_artist["updated_at"] = now

        query = (
            insert(DimAlbumArtist)
            .values(album_artists_to_write)
            .on_conflict_do_nothing(
                index_elements=get_model_pk(DimAlbumArtist),
            )
        )

        await async_execute(etldb_engine, query)


async def store_track_artists_to_etldb(
    etldb_engine: AsyncEngine,
    model: EtlModelType,
    entity: EntityType,
) -> None:
    if isinstance(model, (SpotifyTrack, AppleMusicTrack)) and model.track_artists:
        track_artists_to_write = model.dict(include={"track_artists"})["track_artists"]

        for track_artist in track_artists_to_write:
            now = datetime.now(timezone.utc)
            track_artist["created_at"] = now
            track_artist["updated_at"] = now

        for track_artist in track_artists_to_write:
            query = (
                insert(DimTrackArtist)
                .values(track_artist)
                .on_conflict_do_update(
                    index_elements=get_model_pk(DimTrackArtist),
                    set_={
                        "artist_id": track_artist["artist_id"],
                        "updated_at": datetime.now(timezone.utc),
                    },
                )
            )
            await async_execute(etldb_engine, query)


async def store_fact_playlist_followers(etldb_engine: AsyncEngine, **artist_followers) -> None:
    query = insert(FactPlaylistFollowers).values(artist_followers)
    playlist_query = select(DimPlaylist).where(
        DimPlaylist.playlist_id == artist_followers["playlist_id"],
        DimPlaylist.dsp_id == artist_followers["dsp_id"],
        DimPlaylist.market_id == artist_followers["market_id"],
    )
    playlist_insert_query = insert(DimPlaylist).values(
        playlist_id=artist_followers["playlist_id"],
        dsp_id=artist_followers["dsp_id"],
        market_id=artist_followers["market_id"],
    )
    result = await async_execute(etldb_engine, playlist_query)
    playlist = result.first()

    if not playlist:
        await async_execute(etldb_engine, playlist_insert_query)

    await async_execute(etldb_engine, query)


async def store_data_to_s3(
    s3_client: S3AsyncClient, entity: EntityType, records: list[Row], data: Any
) -> None:
    if not entity.config.store_raw:
        return

    for record in records:
        api_data_for_record = entity.get_record_from_api_response(record, data)
        await store_record_data_to_s3(s3_client, entity, record, api_data_for_record)


@async_retry()
async def store_record_data_to_s3(
    s3_client: S3AsyncClient, entity: EntityType, record: Row, record_data: Any
):
    if not record_data:
        return

    if hasattr(record, "snapshot_id") and record.snapshot_id == record_data.get("snapshot_id"):
        return

    path = get_s3_path_for_file(entity, record)

    with statsd.timed(
        "dapd-api-scraper.s3_write_time",
        tags=[
            f"dsp:{entity.data_source}",
            f"entity:{entity.entity_type}",
        ],
        use_ms=True,
    ):
        await s3_client.put_object(
            path,
            gzip.compress(orjson.dumps(record_data)),
        )


def get_s3_path_for_file(entity: EntityType, record: Row) -> str:
    now = datetime.utcnow()
    return (
        f"ds={entity.data_source}/"
        f"entity={entity.entity_type}/"
        f"storefront={record.storefront}/"
        f"y={now.year}/m={now.month}/d={now.day}/"
        f"h={now.hour}/{record.id}.json"
    )


@async_cache()
async def get_market_id_by_storefront(etldb_engine: AsyncEngine, storefront: str) -> int:
    """
    Returns market_id using storefront as a filter, in case if market with the
    given storefront does not exists will be crated one"""
    query = select(DimMarket).where(DimMarket.market_code == storefront)

    result = await async_execute(etldb_engine, query)
    market = result.first()

    if not market:
        query = (
            insert(DimMarket)
            .values(
                market_code=storefront,
                market_name=storefront,
                created_at=datetime.now(timezone.utc),
            )
            .returning(DimMarket.market_id)
        )
        result = await async_execute(etldb_engine, query)
        market = result.one()

    market_id: int = market.market_id

    return market_id


@async_cache()
async def get_dsp_id_by_data_source(etldb_engine: AsyncEngine, data_source: str) -> int:
    """
    Returns dsp_id using data_source as a filter, in case if market with the
    given storefront does not exists will be crated one"""
    query = select(DimDsp).where(DimDsp.dsp_name == data_source)
    result = await async_execute(etldb_engine, query)
    dsp = result.first()

    if not dsp:
        query = (
            insert(DimMarket)
            .values(
                dsp_name=data_source,
                created_at=datetime.now(timezone.utc),
            )
            .returning(DimDsp.dsp_id)
        )
        result = await async_execute(etldb_engine, query)
        dsp = result.one()

    dsp_id: int = dsp.dsp_id

    return dsp_id


@cache
def get_model_pk(model):
    """Return sqlalchemy model's primary key"""
    return inspect(model).primary_key