import imghdr
import base64
import six
import uuid

from django.conf import settings
from django.core.files.base import ContentFile

from rest_framework import serializers


class Base64ImageField(serializers.ImageField):
    use_url = False

    def to_internal_value(self, data):
        if isinstance(data, six.string_types):
            if 'data:' in data and ';base64,' in data:
                header, data = data.split(';base64,')
                try:
                    decoded_file = base64.b64decode(data)
                except TypeError:
                    self.fail('invalid_image')

                file_name = str(uuid.uuid4())[:12]
                file_extension = self.get_file_extension(
                    file_name, decoded_file)
                complete_file_name = "%s.%s" % (file_name, file_extension,)
                data = ContentFile(decoded_file, name=complete_file_name)
        return super().to_internal_value(data)

    def get_file_extension(self, file_name, decoded_file):
        extension = imghdr.what(file_name, decoded_file)
        extension = "jpg" if extension == "jpeg" else extension
        return extension

    def to_representation(self, value):
        file = "%s%s" % (settings.MEDIA_URL, super().to_representation(value),)

        return file