from __future__ import unicode_literals

from django import forms
from django.forms.models import modelform_factory
from django.utils.module_loading import import_string
from django.utils.translation import gettext_lazy as _

from wagtail.admin import widgets
from wagtail.admin.forms.collections import (
    BaseCollectionMemberForm,
    collection_member_permission_formset_factory,
)

from wagtailmedia.models import Media
from wagtailmedia.permissions import permission_policy as media_permission_policy
from wagtailmedia.settings import wagtailmedia_settings


class BaseMediaForm(BaseCollectionMemberForm):
    class Meta:
        widgets = {
            "tags": widgets.AdminTagWidget,
            "file": forms.FileInput,
            "thumbnail": forms.ClearableFileInput,
        }

    permission_policy = media_permission_policy

    def __init__(self, *args, **kwargs):
        super(BaseMediaForm, self).__init__(*args, **kwargs)

        if self.instance.type == "audio":
            for name in ("width", "height"):
                # these fields might be editable=False so verify before accessing
                if name in self.fields:
                    del self.fields[name]


def get_media_base_form():
    base_form_override = wagtailmedia_settings.MEDIA_FORM_BASE
    if base_form_override:
        base_form = import_string(base_form_override)
    else:
        base_form = BaseMediaForm
    return base_form


def get_media_form(model):
    fields = model.admin_form_fields
    if "collection" not in fields:
        # force addition of the 'collection' field, because leaving it out can
        # cause dubious results when multiple collections exist (e.g adding the
        # media to the root collection where the user may not have permission) -
        # and when only one collection exists, it will get hidden anyway.
        fields = list(fields) + ["collection"]

    return modelform_factory(
        model,
        form=get_media_base_form(),
        fields=fields,
    )


GroupMediaPermissionFormSet = collection_member_permission_formset_factory(
    Media,
    [
        ("add_media", _("Add"), _("Add/edit media you own")),
        ("change_media", _("Edit"), _("Edit any media")),
    ],
    "wagtailmedia/permissions/includes/media_permissions_formset.html",
)

@@ -329,16 +329,43 @@
 erForm,%0A
+    CollectionChoiceField,%0A
     coll
@@ -587,16 +587,587 @@
 tings%0A%0A%0A
+try:%0A    from wagtail.models import Collection%0Aexcept ImportError:%0A    from wagtail.core.models import Collection%0A%0A%0A# Callback to allow us to override the default form field for the collection field%0Adef formfield_for_dbfield(db_field, **kwargs):%0A    if db_field.name == %22collection%22:%0A        return CollectionChoiceField(%0A            label=_(%22Collection%22),%0A            queryset=Collection.objects.all(),%0A            empty_label=None,%0A            **kwargs,%0A        )%0A%0A    # For all other fields, just call its formfield() method.%0A    return db_field.formfield(**kwargs)%0A%0A%0A
 class Ba
@@ -2592,16 +2592,66 @@
 fields,%0A
+        formfield_callback=formfield_for_dbfield,%0A
     )%0A%0A%0A

Update the Media form collection field override to follow core