# Generated by Django 2.0.13 on 2019-02-28 11:28

import json

from urllib.parse import unquote_plus

from django.db import migrations


def remove_signatures(apps, schema_editor):
    Recipe = apps.get_model("recipes", "Recipe")
    Action = apps.get_model("recipes", "Action")
    Signature = apps.get_model("recipes", "Signature")

    for recipe in Recipe.objects.exclude(signature=None):
        sig = recipe.signature
        recipe.signature = None
        recipe.save()
        sig.delete()

    for action in Action.objects.exclude(signature=None):
        sig = action.signature
        action.signature = None
        action.save()
        sig.delete()

    for sig in Signature.objects.all():
        sig.delete()


def add_extension_id(apps, schema_editor):
    Action = apps.get_model("recipes", "Action")
    RecipeRevision = apps.get_model("recipes", "RecipeRevision")
    Extension = apps.get_model("studies", "Extension")

    action = Action.objects.get(name="opt-out-study")
    revisions = RecipeRevision.objects.filter(action_id=action.id)

    for revision in revisions:
        arguments = json.loads(revision.arguments_json)
        url = arguments.get("addonUrl")
        filename = unquote_plus(url.split("/extensions/").pop())
        extension = Extension.objects.get(xpi=f"extensions/{filename}")
        arguments["extensionApiId"] = extension.id
        revision.arguments_json = json.dumps(arguments)
        revision.save()

        # Remove signatures to prompt resigning
        remove_signatures(apps, schema_editor)


def remove_extension_id(apps, schema_editor):
    Action = apps.get_model("recipes", "Action")
    RecipeRevision = apps.get_model("recipes", "RecipeRevision")

    action = Action.objects.get(name="opt-out-study")
    revisions = RecipeRevision.objects.filter(action_id=action.id)

    for revision in revisions:
        arguments = json.loads(revision.arguments_json)
        if "extensionApiId" in arguments:
            arguments.pop("extensionApiId")
        revision.arguments_json = json.dumps(arguments)
        revision.save()

        # Remove signatures to prompt resigning
        remove_signatures(apps, schema_editor)


class Migration(migrations.Migration):

    dependencies = [
        ("recipes", "0013_auto_20181018_2049"),
        ("studies", "0006_extension_hash_algorithm"),
    ]

    operations = [migrations.RunPython(add_extension_id, remove_extension_id)]

Add migration to add extension ID to addon study recipes