import gradio as gr
import os
import requests
import json
import base64
from io import BytesIO
from PIL import Image
from huggingface_hub import login

from css_html_js import custom_css

from about import (
    CITATION_BUTTON_LABEL,
    CITATION_BUTTON_TEXT,
    EVALUATION_QUEUE_TEXT,
    INTRODUCTION_TEXT,
    LLM_BENCHMARKS_TEXT,
    TITLE,
)

myip = "146.152.224.103"
myport=8080
is_spaces = True if "SPACE_ID" in os.environ else False
is_shared_ui = False


def process_image_from_binary(img_stream):
    if img_stream is None:
        print("no image binary")
        return
    image_data = base64.b64decode(img_stream)
    image_bytes = BytesIO(image_data)
    img = Image.open(image_bytes)
    
    return img

def generate_img(concept, prompt, seed, steps):
    print(f"my IP is {myip}, my port is {myport}")
    response = requests.post('http://{}:{}/generate'.format(myip, myport), 
                             json={"concept": concept, "prompt": prompt, "seed": seed, "steps": steps},
                             timeout=(10, 1200))
    print(f"result: {response}")
    image = None
    if response.status_code == 200:
        response_json = response.json()
        print(response_json)
        image = process_image_from_binary(response_json['image'])
    else:
        print(f"Request failed with status code {response.status_code}")
    
    return image

with gr.Blocks() as demo:
    gr.HTML(TITLE)
    gr.Markdown(INTRODUCTION_TEXT, elem_classes="markdown-text")
    with gr.Row() as advlearn:
        with gr.Column():
            # gr.Markdown("Please upload your model id.")
            drop_text = gr.Dropdown(["Object-Church", "Object-Parachute", "Object-Garbage_Truck",
                                "Style-VanGogh","Concept-Nudity", "None"], 
                                   label="AdvUnlearn Text Encoder")
        with gr.Column():
            text_input = gr.Textbox(label="Prompt")
        
    with gr.Row():
        with gr.Column():
            with gr.Row():
                seed = gr.Textbox(label="seed", value=666)
            with gr.Row():
                steps = gr.Textbox(label="num_steps", value=100)
            with gr.Row():
                start_button = gr.Button("AdvUnlearn",size='lg')
        with gr.Column(min_width=512):
            result_img = gr.Image(label="Image Gnerated by AdvUnlearn",width=512,show_share_button=False,show_download_button=False)

    start_button.click(fn=generate_img, inputs=[drop_text, text_input, seed, steps], outputs=result_img, api_name="generate")
            
demo.launch()