import os.path as op
from io import BytesIO

import boto3
from botocore.exceptions import ClientError
from flask_admin._compat import urljoin
from flask_admin.form.upload import FileUploadField, ImageUploadInput, thumbgen_filename
from werkzeug.datastructures import FileStorage
from wtforms import ValidationError

from src.utils import get_s3_image_key_from_url


try:
    from PIL import Image, ImageOps
except ImportError:
    Image = None
    ImageOps = None


class S3FileUploadField(FileUploadField):
    """
    Inherits from flask-admin FileUploadField, to allow file uploading to Amazon S3.
    """
    def __init__(
        self, label=None, validators=None, bucket_name=None, acl="public-read", static_root_parent=None, **kwargs
    ):
        super(S3FileUploadField, self).__init__(label, validators, **kwargs)
        self.bucket_name = bucket_name
        self.acl = acl
        self.static_root_parent = static_root_parent

    def populate_obj(self, obj, name):
        if self.data and isinstance(self.data, FileStorage) and self.data.filename:
            filename = self.generate_name(obj, self.data)
            temp_file = BytesIO()
            self.data.save(temp_file)
            filename = self._save_file(temp_file, filename)
            self.data.filename = filename

            setattr(obj, name, filename)

    def _delete_file(self, *args, **kwargs):
        pass

    def _save_file(self, temp_file, filename):
        s3_client = boto3.client("s3")
        temp_file.seek(0)
        s3_client.upload_fileobj(temp_file, self.bucket_name, filename, ExtraArgs={"ACL": self.acl})
        return filename


class S3ImageUploadInput(ImageUploadInput):
    """
    Inherits from flask-admin ImageUploadInput, to render images uploaded to Amazon S3.
    """

    def get_url(self, field):
        if op.isfile(op.join(field.base_path, field.data)):
            return super(S3ImageUploadInput, self).get_url(field)

        filename = field.data

        if field.url_relative_path:
            filename = urljoin(field.url_relative_path, filename)

        return filename


class S3ImageUploadField(S3FileUploadField):
    """
    Revised version of flask-admin ImageUploadField, to allow image uploading to Amazon S3.
    """

    widget = S3ImageUploadInput()

    keep_image_formats = ("PNG",)

    def __init__(
        self,
        label=None,
        validators=None,
        max_size=None,
        thumbgen=None,
        thumbnail_size=None,
        url_relative_path=None,
        endpoint="static",
        **kwargs,
    ):
        # Check if PIL is installed
        if Image is None:
            raise ImportError("PIL library was not found")

        self.max_size = max_size
        self.thumbnail_fn = thumbgen or thumbgen_filename
        self.thumbnail_size = thumbnail_size
        self.endpoint = endpoint
        self.image = None
        self.url_relative_path = url_relative_path

        if not ("allowed_extensions" in kwargs) or not kwargs["allowed_extensions"]:
            kwargs["allowed_extensions"] = ("gif", "jpg", "jpeg", "png", "tiff")

        super(S3ImageUploadField, self).__init__(label, validators, **kwargs)

    def pre_validate(self, form):
        super(S3ImageUploadField, self).pre_validate(form)

        if self.data and isinstance(self.data, FileStorage) and self.data.filename:
            try:
                self.image = Image.open(self.data)
            except Exception as e:
                raise ValidationError("Invalid image: %s" % e)

    def _delete_file(self, *args, **kwargs):
        pass

    def _save_file(self, temp_file, filename):
        filename, format = self._get_save_format(filename, self.image)

        if self.image:
            temp_file = BytesIO()
            self._save_image(self.image, temp_file, format)

        super(S3ImageUploadField, self)._save_file(temp_file, filename)
        return filename

    def _save_image(self, image, temp_file, format="JPEG"):
        if image.mode not in ("RGB", "RGBA"):
            image = image.convert("RGBA")

        image.save(temp_file, format)

    def _get_save_format(self, filename, image):
        if image.format not in self.keep_image_formats:
            name, ext = op.splitext(filename)
            filename = "%s.jpg" % name
            return filename, "JPEG"

        return filename, image.format


def delete_image_by_url(image_url: str, bucket_name: str):

    img_key = get_s3_image_key_from_url(image_url)
    s3 = boto3.resource("s3")

    try:
        s3.Object(bucket_name, img_key).delete()
    except ClientError:
        pass