import asyncio

from abstract import AbstractMigration


class CustomMigration(AbstractMigration):
    def filter(self, manifest_metadata: dict) -> bool:
        return manifest_metadata['environment'] == "qa"

    def migrate(self, manifest: dict) -> dict:
        if 'type' in manifest and manifest['type'] == 'deployment':
            return manifest

        (manifest.setdefault('flink', {})
         .setdefault('flinkDeployment', {})
         .setdefault('flinkConfiguration', {})
         .setdefault('flinkVersion', "v1_17"))

        try:
            del (manifest['flink']['flinkDeployment']['taskManager']['replicas'])
        except KeyError:
            pass

        manifest['flink']['flinkDeployment']['flinkConfiguration']['flinkVersion'] = "v1_17"

        return manifest


async def main():
    migration = CustomMigration(
        "https://deployment-portal-backend.np.navi-tech.in",
        "<put your x-auth-token here>",
        # Dry run is enabled by default. Uncomment this line to disable dry run
        # dry_run=False
    )
    await migration.run()


if __name__ == "__main__":
    asyncio.run(main())