import time
import uuid

import gradio as gr
import numpy as np
from PIL import Image

from .generation import call_generation
from .util import upload_np_2_oss


def get_image_url(request_id, pil_image, suffix):
    np_image = np.array(pil_image.convert("RGB"))
    image_url = upload_np_2_oss(np_image, request_id + suffix)
    return image_url


def is_no_image_placeholder(pil_image):
    from PIL import Image, ImageChops
    no_image_placeholder = Image.open("assets/No-Image-Placeholder.png").convert("RGB")
    diff = ImageChops.difference(pil_image, no_image_placeholder)
    if diff.getbbox() is None:
        return True
    else:
        return False


def text_to_single_id_generation_process(
        pil_faceid=None, pil_supp_faceids=None,
        pil_mix_faceid_1=None, mix_scale_1=0.0,
        pil_mix_faceid_2=None, mix_scale_2=0.0,
        faceid_scale=0.0, face_structure_scale=0.0,
        prompt="", negative_prompt="",
        pil_style=None, style_scale=0.0,
        seed=-1, image_resolution="512x512", use_sr=True,
):
    request_id = time.strftime('%Y%m%d-', time.localtime(time.time())) + str(uuid.uuid4())

    if prompt == "":
        raise gr.Error("Please enter the prompt")

    if pil_style and is_no_image_placeholder(pil_style):
        pil_style = None

    faceid_url = []
    if pil_faceid:
        faceid_url.append(get_image_url(request_id, pil_faceid, suffix="_faceid.png"))
    if pil_supp_faceids and len(pil_supp_faceids) > 0:
        for idx, pil_supp_faceid in enumerate(pil_supp_faceids):
            pil_supp_faceid = Image.open(pil_supp_faceid)
            faceid_url.append(get_image_url(request_id, pil_supp_faceid, suffix=f"_supp_faceid_{idx}.png"))

    mix_face_url = []
    mix_scale = []
    if pil_mix_faceid_1 is not None:
        mix_face_url.append(
            get_image_url(request_id, pil_mix_faceid_1, suffix=f"_mix_faceid_1_{mix_scale_1:.2f}.png")
        )
        mix_scale.append(mix_scale_1)

    if pil_mix_faceid_2 is not None:
        mix_face_url.append(
            get_image_url(request_id, pil_mix_faceid_2, suffix=f"_mix_faceid_2_{mix_scale_2:.2f}.png")
        )
        mix_scale.append(mix_scale_2)

    # The outer parenthesis represent the ID, and the inner parenthesis represent the mix face image for each ID.
    faceid_image_url, mix_faceid_image_url, mix_faceid_scale = [], [], []
    if len(faceid_url) > 0:
        faceid_image_url.append(faceid_url)
        mix_faceid_image_url.append(mix_face_url)
        mix_faceid_scale.append(mix_scale)

    if pil_style is None:
        style_image_url = ""
    else:
        style_image_url = get_image_url(request_id, pil_style, suffix="_style.png")

    data = dict(
        input=dict(
            faceid_image_url=faceid_image_url,
            mix_faceid_image_url=mix_faceid_image_url,
            mix_faceid_scale=mix_faceid_scale,
            style_image_url=style_image_url,
            prompt=prompt,
            negative_prompt=negative_prompt
        ),
        parameters=dict(
            image_resolution=image_resolution,
            faceid_scale=faceid_scale * 10.,
            face_structure_scale=face_structure_scale * 10.,
            style_scale=style_scale * 10.,
            use_sr=int(use_sr),
            seed=seed,
            mode="text-to-image",
        )
    )

    res = call_generation(data)

    return res


def text_to_multi_id_generation_process(
        pil_faceid_1st=None, pil_supp_faceids_1st=None,
        pil_mix_faceid_1_1st=None, mix_scale_1_1st=0.0,
        pil_mix_faceid_2_1st=None, mix_scale_2_1st=0.0,
        pil_faceid_2nd=None, pil_supp_faceids_2nd=None,
        pil_mix_faceid_1_2nd=None, mix_scale_1_2nd=0.0,
        pil_mix_faceid_2_2nd=None, mix_scale_2_2nd=0.0,
        faceid_scale=0.0, face_structure_scale=0.0,
        prompt="", negative_prompt="",
        pil_style=None, style_scale=0.0,
        seed=-1, image_resolution="512x512", use_sr=True,
):
    request_id = time.strftime('%Y%m%d-', time.localtime(time.time())) + str(uuid.uuid4())

    if prompt == "":
        raise gr.Error("Please enter the prompt")

    if pil_style and is_no_image_placeholder(pil_style):
        pil_style = None

    first_faceid_url = []
    if pil_faceid_1st:
        first_faceid_url.append(get_image_url(request_id, pil_faceid_1st, suffix="_faceid_1st.png"))
    if pil_supp_faceids_1st and len(pil_supp_faceids_1st) > 0:
        for idx, pil_supp_faceid_1st in enumerate(pil_supp_faceids_1st):
            pil_supp_faceid_1st = Image.open(pil_supp_faceid_1st)
            first_faceid_url.append(get_image_url(request_id, pil_supp_faceid_1st, suffix=f"_faceid_1st_{idx}.png"))

    second_faceid_url = []
    if pil_faceid_2nd:
        second_faceid_url.append(get_image_url(request_id, pil_faceid_2nd, suffix="_faceid_2nd.png"))
    if pil_supp_faceids_2nd and len(pil_supp_faceids_2nd) > 0:
        for idx, pil_supp_faceid_2nd in enumerate(pil_supp_faceids_2nd):
            pil_supp_faceid_2nd = Image.open(pil_supp_faceid_2nd)
            second_faceid_url.append(get_image_url(request_id, pil_supp_faceid_2nd, suffix=f"_faceid_2nd_{idx}.png"))

    first_mix_face_url, first_mix_scale = [], []
    if pil_mix_faceid_1_1st:
        first_mix_face_url.append(
            get_image_url(request_id, pil_mix_faceid_1_1st, suffix=f"_mix_faceid_1_1st_{mix_scale_1_1st:.2f}.png")
        )
        first_mix_scale.append(mix_scale_1_1st)

    if pil_mix_faceid_2_1st is not None:
        first_mix_face_url.append(
            get_image_url(request_id, pil_mix_faceid_2_1st, suffix=f"_mix_faceid_2_1st_{mix_scale_2_1st:.2f}.png")
        )
        first_mix_scale.append(mix_scale_2_1st)

    second_mix_face_url, second_mix_scale = [], []
    if pil_mix_faceid_1_2nd:
        second_mix_face_url.append(
            get_image_url(request_id, pil_mix_faceid_1_2nd, suffix=f"_mix_faceid_1_2nd_{mix_scale_1_2nd:.2f}.png")
        )
        second_mix_scale.append(mix_scale_1_2nd)

    if pil_mix_faceid_2_2nd is not None:
        second_mix_face_url.append(
            get_image_url(request_id, pil_mix_faceid_2_2nd, suffix=f"_mix_faceid_2_2nd_{mix_scale_2_2nd:.2f}.png")
        )
        second_mix_scale.append(mix_scale_2_2nd)

    # The outer parenthesis represent the ID, and the inner parenthesis represent the mix face image for each ID.
    faceid_image_url, mix_faceid_image_url, mix_faceid_scale = [], [], []
    if len(first_faceid_url) > 0:
        faceid_image_url.append(first_faceid_url)
        mix_faceid_image_url.append(first_mix_face_url)
        mix_faceid_scale.append(first_mix_scale)
    if len(second_faceid_url) > 0:
        faceid_image_url.append(second_faceid_url)
        mix_faceid_image_url.append(second_mix_face_url)
        mix_faceid_scale.append(second_mix_scale)

    if pil_style is None:
        style_image_url = ""
    else:
        style_image_url = get_image_url(request_id, pil_style, suffix="_style.png")

    data = dict(
        input=dict(
            faceid_image_url=faceid_image_url,
            mix_faceid_image_url=mix_faceid_image_url,
            mix_faceid_scale=mix_faceid_scale,
            style_image_url=style_image_url,
            prompt=prompt,
            negative_prompt=negative_prompt
        ),
        parameters=dict(
            image_resolution=image_resolution,
            faceid_scale=faceid_scale * 10.,
            face_structure_scale=face_structure_scale * 10.,
            style_scale=style_scale * 10.,
            use_sr=int(use_sr),
            seed=seed,
            mode="text-to-image",
        )
    )

    res = call_generation(data)

    return res


def image_to_single_id_generation_process(
        pil_faceid=None, pil_supp_faceids=None,
        pil_mix_faceid_1=None, mix_scale_1=0.0,
        pil_mix_faceid_2=None, mix_scale_2=0.0,
        faceid_scale=0.0, face_structure_scale=0.0,
        pil_style=None, style_scale=1.0,
        seed=-1, image_resolution="768x512", use_sr=True,
):
    request_id = time.strftime('%Y%m%d-', time.localtime(time.time())) + str(uuid.uuid4())

    if pil_faceid is None:
        raise gr.Error("Please upload an ID image")

    if pil_style is None:
        raise gr.Error("Please upload a reference image")

    faceid_url = []
    if pil_faceid:
        faceid_url.append(get_image_url(request_id, pil_faceid, suffix="_faceid.png"))
    if pil_supp_faceids and len(pil_supp_faceids) > 0:
        for idx, pil_supp_faceid in enumerate(pil_supp_faceids):
            pil_supp_faceid = Image.open(pil_supp_faceid)
            faceid_url.append(get_image_url(request_id, pil_supp_faceid, suffix=f"_supp_faceid_{idx}.png"))

    mix_face_url = []
    mix_scale = []
    if pil_mix_faceid_1 is not None:
        mix_face_url.append(
            get_image_url(request_id, pil_mix_faceid_1, suffix=f"_mix_faceid_1_{mix_scale_1:.2f}.png")
        )
        mix_scale.append(mix_scale_1)

    if pil_mix_faceid_2 is not None:
        mix_face_url.append(
            get_image_url(request_id, pil_mix_faceid_2, suffix=f"_mix_faceid_2_{mix_scale_2:.2f}.png")
        )
        mix_scale.append(mix_scale_2)

    # The outer parenthesis represent the ID, and the inner parenthesis represent the mix face image for each ID.
    faceid_image_url, mix_faceid_image_url, mix_faceid_scale = [], [], []
    if len(faceid_url) > 0:
        faceid_image_url.append(faceid_url)
        mix_faceid_image_url.append(mix_face_url)
        mix_faceid_scale.append(mix_scale)

    if pil_style is None:
        style_image_url = ""
    else:
        style_image_url = get_image_url(request_id, pil_style, suffix="_style.png")

    data = dict(
        input=dict(
            faceid_image_url=faceid_image_url,
            mix_faceid_image_url=mix_faceid_image_url,
            mix_faceid_scale=mix_faceid_scale,
            style_image_url=style_image_url,
        ),
        parameters=dict(
            image_resolution=image_resolution,
            faceid_scale=faceid_scale * 10.,
            face_structure_scale=face_structure_scale * 10.,
            style_scale=style_scale * 10.,
            use_sr=int(use_sr),
            seed=seed,
            mode="image-to-image",
        )
    )

    res = call_generation(data)

    return res