from typing import Dict, Any

from mongodb_migrations.base import BaseMigration


class Migration(BaseMigration):
    def upgrade(self):
        index_information_users: Dict[str, Any] = self.db.users.index_information()
        index_information_events: Dict[str, Any] = self.db.events.index_information()
        index_information_stages: Dict[str, Any] = self.db.stages.index_information()

        # Update users collection
        if "user_id" in index_information_users:
            self.db.users.drop_index("user_id")

        if "user_id-guild_id" not in index_information_users:
            self.db.users.create_index(["id", "guild_id"], name="user_id-guild_id", unique=True)

        # Update events collection
        if "guild_id" in index_information_events:
            self.db.events.drop_index("guild_id")

        if "event_name-guild_id" not in index_information_events:
            self.db.events.create_index(["name", "guild_id"], name="event_name-guild_id", unique=False)

        # Update stages collection
        if "event_id-and-guild_id" in index_information_stages:
            self.db.stages.drop_index("event_id-and-guild_id")

        if "event_id-guild_id" not in index_information_stages:
            self.db.stages.create_index(["event_id", "guild_id"], name="event_id-guild_id", unique=False)

    def downgrade(self):
        index_information_users: Dict[str, Any] = self.db.users.index_information()
        index_information_events: Dict[str, Any] = self.db.events.index_information()
        index_information_stages: Dict[str, Any] = self.db.stages.index_information()

        # Update users collection
        if "user_id-guild_id" in index_information_users:
            self.db.users.drop_index("user_id-guild_id")

        if "user_id" not in index_information_users:
            self.db.users.create_index("id", name="user_id", unique=True)

        # Update events collection
        if "event_name-guild_id" in index_information_events:
            self.db.events.drop_index("event_name-guild_id")

        if "guild_id" not in index_information_events:
            self.db.events.create_index("guild_id", name="guild_id", unique=False)

        # Update stages collection
        if "event_id-guild_id" in index_information_stages:
            self.db.stages.drop_index("event_id-guild_id")

        if "event_id-and-guild_id" not in index_information_stages:
            self.db.stages.create_index(
                ["event_id", "guild_id"], name="event_id-and-guild_id", unique=False
            )