from mongodb_migrations.base import BaseMigration


class Migration(BaseMigration):
    name_template = "ind_{}_created_at"

    collections_list = (
        "spotify_track",
        "spotify_track_isrc",
        "spotify_playlist",
        "spotify_album",
        "spotify_artist",
        "spotify_user",
        "spotify_search",
        "spotify_album_tracks",
        "apple_song",
        "apple_song_isrc",
        "apple_playlist",
        "apple_album",
        "apple_search",
        "spotify_playlist_tracks",
        "apple_station",
        "spotify_cache",
        "apple_cache"
    )

    def set_index_ttl(self, timeout: int):
        for collection_name in self.collections_list:
            index_name = self.name_template.format(collection_name)
            self.db[collection_name].drop_index(index_name)
            self.db[collection_name].create_index(
                "created_at", name=index_name, expireAfterSeconds=timeout, background=True,
            )

    def upgrade(self):
        # ttl index 1 day
        self.set_index_ttl(1 * 24 * 60 * 60)

    def downgrade(self):
        # ttl index 30 days
        self.set_index_ttl(30 * 24 * 60 * 60)