import gradio as gr
import hmac
import hashlib
import time
import os
import requests
from io import BytesIO
from PIL import Image
import uuid

example_path = os.path.join(os.path.dirname(__file__), 'assets')
clothing_list = os.listdir(os.path.join(example_path, "clothing"))
clothing_list_path = [os.path.join(example_path, "clothing", clothing) for clothing in clothing_list]
human_list = os.listdir(os.path.join(example_path, "human"))
human_list_path = [os.path.join(example_path, "human", human) for human in human_list]

base_url = os.getenv('base_url')
upload_image_url = os.getenv('upload_image_url')
upload_background_image_url = os.getenv('upload_background_image_url')
create_save_task_url = os.getenv('create_save_task_url')
execute_task_url = os.getenv('execute_task_url')
query_task_url = os.getenv('query_task_url')
secret_key = os.getenv('secret_key')
agent_version = os.getenv('agent_version')
agent_name = os.getenv('agent_name')


def parse_response(response):
    data = {}
    msg = ''
    if response.status_code == 200:
        try:
            datas = response.json()
            if datas:
                data = datas.get("data")
                if not data:
                    msg = datas.get("msg")
                    if not msg:
                        msg = "Field error."
            else:
                msg = "The parsing result is empty."
        except Exception as e:
            msg = f"parse error: {repr(e)}."
    else:
        msg = f'request error.'
    return data, msg


def generate_signature(key, did, timestamp):
    data = f"{did}:{timestamp}"
    h = hmac.new(key.encode(), data.encode(), hashlib.sha256)
    return h.hexdigest()


def url_to_image(url):
    headers = {
        "user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36"
    }
    try:
        response = requests.get(url, headers=headers, timeout=30)
    except:
        return None
    if response.status_code == 200:
        img = Image.open(BytesIO(response.content))
        return img
    return None


def start_task(task_id, timestamp, signature, did):
    headers = {
        'Did': did,
        'X-Timestamp': timestamp,
        'X-Signature': signature
    }
    data = {
        "agentVersion": agent_version,
        "agentName": agent_name,
        "taskId": task_id,
        "runFreeAsFallback": False
    }
    response = requests.post(base_url + execute_task_url, json=data, headers=headers)
    data, msg = parse_response(response)
    return data, msg


def create_task(image_url, timestamp, signature, did):
    headers = {
        'Did': did,
        'X-Timestamp': timestamp,
        'X-Signature': signature
    }
    data = {
        "agentVersion": agent_version,
        "agentName": agent_name,
        "image": image_url
    }
    response = requests.post(base_url + create_save_task_url, json=data, headers=headers)
    data, msg = parse_response(response)
    return data, msg


def save_task(image_url, timestamp, signature, show_image, task_id, location_id, did):
    headers = {
        'Did': did,
        'X-Timestamp': timestamp,
        'X-Signature': signature
    }
    data = {
        "agentVersion": agent_version,
        "agentName": agent_name,
        "image": image_url,
        "showImage": show_image,
        "taskId": task_id,
        "locationId": location_id,
    }
    response = requests.post(base_url + create_save_task_url, json=data, headers=headers)
    data, msg = parse_response(response)
    return data, msg


def query_task(task_id, execution_id, timestamp, signature, did):
    headers = {
        'Did': did,
        'X-Timestamp': timestamp,
        'X-Signature': signature
    }
    data = {
        "agentVersion": agent_version,
        "agentName": agent_name,
        "taskId": task_id,
        "executionId": execution_id,
    }
    response = requests.post(base_url + query_task_url, json=data, headers=headers)
    data, msg = parse_response(response)
    return data, msg


def upload_image(image, signature, timestamp, upload_type, did):
    if image is None:
        return None
    if upload_type == 'image':
        upload_url = upload_image_url
    else:
        upload_url = upload_background_image_url
    image_format = image.format if image.format else "PNG"
    mime_type = f"image/{image_format.lower()}"
    with BytesIO() as m_img:
        image.save(m_img, format=image_format)
        m_img.seek(0)
        files = {'image': (f"main_image.{image_format.lower()}", m_img, mime_type)}
        headers = {
            'Did': did,
            'X-Timestamp': timestamp,
            'X-Signature': signature
        }
        response = requests.post(base_url + upload_url, files=files, headers=headers)
        data, msg = parse_response(response)
        return data, msg


def load_description(file_path):
    with open(file_path, 'r', encoding='utf-8') as file:
        content = file.read()
    return content


def generate_image(main_image, background_image, did):
    if not did:
        did = str(uuid.uuid4())
    if main_image is None or background_image is None:
        m = "Please upload both the main image and the background reference image before generating."
        return gr.Warning(m), did

    timestamp = str(int(time.time()))
    signature = generate_signature(
        key=secret_key,
        did=did,
        timestamp=timestamp
    )
    upload_image_data, upload_image_msg = upload_image(
        image=main_image,
        signature=signature,
        timestamp=timestamp,
        upload_type='image',
        did=did
    )
    if not upload_image_data:
        return gr.Warning(upload_image_msg), did
    image_url = upload_image_data.get("image")
    if not image_url:
        m = 'Upload image failed.'
        return gr.Warning(m), did
    create_task_data, create_task_msg = create_task(
        image_url=image_url,
        timestamp=timestamp,
        signature=signature,
        did=did
    )
    if not create_task_data:
        return gr.Warning(create_task_msg), did
    task_id = create_task_data.get("taskId")
    show_image = create_task_data.get("showImage")
    if not task_id or not show_image:
        m = 'Create task failed.'
        return gr.Warning(m), did
    upload_image_data, upload_image_msg = upload_image(
        image=background_image,
        signature=signature,
        timestamp=timestamp,
        upload_type='background_image',
        did=did
    )
    if not upload_image_data:
        return gr.Warning(upload_image_msg), did

    save_task_data, save_task_msg = save_task(
        image_url=image_url,
        timestamp=timestamp,
        signature=signature,
        show_image=show_image,
        task_id=task_id,
        location_id=upload_image_data,
        did=did
    )
    if not save_task_data:
        return gr.Warning(save_task_msg), did
    save_task_id = save_task_data.get("taskId")
    save_show_image = save_task_data.get("showImage")
    if not save_task_id or not save_show_image:
        return gr.Warning('Save task failed'), did
    start_task_data, start_task_msg = start_task(
        task_id=save_task_id,
        timestamp=timestamp,
        signature=signature,
        did=did
    )
    if not start_task_data:
        return gr.Warning(start_task_msg), did
    execution_id = start_task_data.get("executionId")
    if not execution_id:
        m = "The task failed to start."
        return gr.Warning(m), did
    start_time = int(time.time())
    while True:
        m = "Query task failed."
        query_task_data, query_task_msg = query_task(
            task_id=save_task_id,
            execution_id=execution_id,
            timestamp=timestamp,
            signature=signature,
            did=did
        )
        if not query_task_data:
            return gr.Warning(query_task_msg), did
        executions = query_task_data.get("executions")
        if not executions:
            return gr.Warning(m), did
        results = executions[0].get("result")
        if not results:
            return gr.Warning(m), did
        status = results[0].get("status")
        if status == "Failed":
            return gr.Warning(m), did
        elif status == "Success" or status == "Blocked":
            img = results[0].get("image")
            if img and str(img).strip() != "":
                return url_to_image(img), did
        end_time = int(time.time())
        if end_time - start_time > 3600:
            m = 'Query task timeout.'
            return gr.Warning(m), did
        time.sleep(2)


def preprocess_image(main_image):
    return main_image


def preprocess_background_image(background_image):
    return background_image


css = """  
    .image-container img {  
        max-height: 500px;  
        width: auto;  
    }  
    .hide-buttons .source-selection {  
        display: none;  
    }  
    #example-images .gallery {
        display: flex;  
        flex-wrap: wrap;  
    }
    #example-images .gallery-item .container{
        width: 100%;
        max-width: 100%;
    }
    #example-images .gallery-item {
        flex: 0 0 32%; 
        max-width: 32%; 
        box-sizing: border-box; 
        display: flex;  
        text-align: center;
        justify-content: center;
        height: 200px;
    }
"""

with gr.Blocks(css=css) as WeShop:
    current_did = gr.State(value='')
    gr.HTML(load_description("assets/title.md"))
    with gr.Row():
        with gr.Column():
            gr.Markdown("#### Step 1: Upload a main image")
            main_image_input = gr.Image(
                type="pil",
                label="Main Image",
                elem_classes=["image-container", "hide-buttons"]
            )
            clothing_example = gr.Examples(
                inputs=main_image_input,
                examples_per_page=12,
                examples=clothing_list_path,
                elem_id="example-images",
                outputs=main_image_input
            )

        with gr.Column():
            gr.Markdown("#### Step 2: Upload a background reference image")
            background_image_input = gr.Image(
                type="pil",
                label="Background reference image",
                elem_classes=["image-container", "hide-buttons"]
            )
            human_example = gr.Examples(
                inputs=background_image_input,
                examples_per_page=12,
                examples=human_list_path,
                elem_id="example-images",
                outputs=background_image_input
            )
        with gr.Column():
            with gr.Row():
                with gr.Column():
                    gr.Markdown("#### Step 3: Press 'Generate' to get the result")
                    output = gr.Image(
                        label="Result",
                        elem_classes=["image-container", "hide-buttons"],
                        interactive=False
                    )
            with gr.Row():
                submit_button = gr.Button("Generate")
                submit_button.click(
                    fn=generate_image,
                    inputs=[main_image_input, background_image_input, current_did],
                    outputs=[output, current_did],
                    concurrency_limit=None
                )
            with gr.Column():
                show_case = gr.Examples(
                    examples=[
                        ["assets/examples/result_02_01.png", "assets/examples/result_02_02.png",
                         "assets/examples/result_02_03.png"],
                        ["assets/examples/result_03_01.png", "assets/examples/result_03_02.png",
                         "assets/examples/result_03_03.png"],
                        ["assets/examples/result_04_01.png", "assets/examples/result_04_02.png",
                         "assets/examples/result_04_03.png"],
                        ["assets/examples/result_05_01.png", "assets/examples/result_05_02.png",
                         "assets/examples/result_05_03.png"],
                        ["assets/examples/result_06_01.png", "assets/examples/result_06_02.png",
                         "assets/examples/result_06_03.png"],
                    ],
                    inputs=[main_image_input, background_image_input, output],
                )
    main_image_input.upload(
        fn=preprocess_image,
        inputs=[main_image_input],
        outputs=main_image_input
    )
    background_image_input.upload(
        fn=preprocess_background_image,
        inputs=[background_image_input],
        outputs=background_image_input
    )
WeShop.queue(api_open=False).launch(show_api=False)