import gradio as gr
import os
import requests
import json
from huggingface_hub import login


# myip = os.environ["34.221.7.12"]
# myport = os.environ["80"]
myip = "34.221.7.12"
myport=80

is_spaces = True if "SPACE_ID" in os.environ else False

is_shared_ui = False

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,
)


def excute_udiff(diffusion_model_id, concept, step):
    print(f"my IP is {myip}, my port is {myport}")
    print(f"my input is diffusion_model_id: {diffusion_model_id}, concept: {concept}, attacker: {attacker}")
    result = requests.post('http://{}:{}/udiff'.format(myip, myport), json={"diffusion_model_id": diffusion_model_id, "concept": concept, "step": step})
    result = result.text[1:-1]
    
    return result


css = '''
    .instruction{position: absolute; top: 0;right: 0;margin-top: 0px !important}
    .arrow{position: absolute;top: 0;right: -110px;margin-top: -8px !important}
    #component-4, #component-3, #component-10{min-height: 0}
    .duplicate-button img{margin: 0}
    #img_1, #img_2, #img_3, #img_4{height:15rem}
    #mdStyle{font-size: 0.7rem}
    #titleCenter {text-align:center}
'''


with gr.Blocks(css=custom_css) as demo:
    gr.HTML(TITLE)
    gr.Markdown(INTRODUCTION_TEXT, elem_classes="markdown-text")

#     gr.Markdown("# Demo of UnlearnDiffAtk.")
#     gr.Markdown("### UnlearnDiffAtk is an effective and efficient adversarial prompt generation approach for unlearned diffusion models(DMs).")
# #     gr.Markdown("####For more details, please visit the [project](https://www.optml-group.com/posts/mu_attack), 
# # check the [code](https://github.com/OPTML-Group/Diffusion-MU-Attack), and read the [paper](https://arxiv.org/abs/2310.11868).")
#     gr.Markdown("### Please notice that the process may take a long time, but the results will be saved. You can try it later if it waits for too long.")
    

    with gr.Row() as udiff:
        with gr.Row():
            drop = gr.Dropdown(["Object-Church", "Object-Parachute", "Object-Garbage_truck","Object-Tench","Style-Van Gogh",
                               "Concept-Nudity", "Concept-Violence", "Concept-Illegal_Activity", "None"], 
                               label="Unlearning undesirable")
        with gr.Column():
            # gr.Markdown("Please upload your model id.")
            drop_model = gr.Dropdown(["Erased Stable Diffusion(ESD)", "Forget-me-not(FMN)", "Ablating concepts(AC)","Unified Concept Editing(UCE)", "(Safe Latent Diffusion)SLD"], 
                               label="Unlearned DMs")
            # diffusion_model_T = gr.Textbox(label='diffusion_model_id')
            # concept = gr.Textbox(label='concept')
            # attacker = gr.Textbox(label='attacker')

            # start_button = gr.Button("Attack!")

        with gr.Column():
             shown_columns_step = gr.Slider(
                            0, 100, value=40, 
                            step=1, label="Attack Steps", info="Choose between 0 and 100",
                            interactive=True,)
    with gr.Row() as attack:
        with gr.Column(min_width=260):
            text_input = gr.Textbox(label="Input Prompt")
            img1 = gr.Image("images/cheetah.jpg",label="Image Generated by Input Prompt",width=260,show_share_button=False,show_download_button=False)
        with gr.Column():
            start_button = gr.Button("UnlearnDiffAtk!",size='lg')
        with gr.Column(min_width=260):
            text_ouput = gr.Textbox(label="Prompt Genetated by UnlearnDiffAtk")
            img2 = gr.Image("images/cheetah.jpg",label="Image Gnerated by Prompt of UnlearnDiffAtk",width=260,show_share_button=False,show_download_button=False)
            
    # with gr.Column():
    #     gr.Examples(examples=[
    #         ["CompVis/stable-diffusion-v1-4", "nudity", "text_grad"]
    #     ], inputs=[diffusion_model_id, concept, attacker])

    start_button.click(fn=excute_udiff, inputs=[drop_model, drop, shown_columns_step], outputs=[text_ouput], api_name="udiff")


# demo.queue(default_enabled=False, api_open=False, max_size=5).launch(debug=True, show_api=False)
demo.queue().launch(server_name='0.0.0.0',share=True)

# with gr.Blocks() as demo:
#     with gr.Row():
#         prompt = gr.Textbox(label='Input Prompt')
#     with gr.Row():
#         shown_columns_1 = gr.CheckboxGroup(
#             choices=["Church","Parachute","Tench", "Garbage Truck"],
#             label="Undersirable Objects",
#             elem_id="column-object",
#             interactive=True,
#         )
#     with gr.Row():
#         shown_columns_2 = gr.CheckboxGroup(
#             choices=["Van Gogh"],
#             label="Undersirable Styles",
#             elem_id="column-style",
#             interactive=True,
#         )
#     with gr.Row():
#         shown_columns_3 = gr.CheckboxGroup(
#             choices=["Violence","Illegal Activity","Nudity"],
#             label="Undersirable Concepts (Outputs that may be offensive in nature)",
#             elem_id="column-select",
#             interactive=True,
#         )
#     with gr.Row():
#         with gr.Column(scale=1, min_width=300):
#             img1 = gr.Image("images/cheetah.jpg",label="Unlearning")
#         with gr.Column(scale=1, min_width=300):
#             img2 = gr.Image("images/cheetah.jpg",label="Attacking")

#     with gr.Row():
#             # gr.Markdown("Please upload your model id.")
#             diffusion_model_id = gr.Textbox(label='diffusion_model_id')
#             shown_columns_4 = gr.Slider(
#                 1, 100, value=40, 
#                 step=1, label="Attacking Steps", info="Choose between 1 and 100",
#                 interactive=True,)

#             # concept = gr.Textbox(label='concept')
#             attacker = gr.Textbox(label='attacker')

#             start_button = gr.Button("Attack!")


# demo.launch()