import gradio as gr
import openai

from openai import OpenAI
# from google.colab import userdata
import os

openai.api_key = "sk-nF3vcQD8BwC7dt6taOyQFH9O36rUBGL5dZd1PieaplT3BlbkFJNKltEcbwN7KItynOiRmFC1VECGPjas9DV7jKcZBdYA"


# 자동 프롬프트 생성
def novel_keyword(nobel_input):
    system_prompt = "당신은 주어지는 내용에 어울리는 이미지를 추천하는 합니다. 이미지를 묘사하는 키워드를 영어로, 콤마로 알려주세요"
    client = OpenAI()
    completion = client.chat.completions.create(
        model="gpt-3.5-turbo",
        temperature=0,
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": nobel_input}
        ])
    return completion.choices[0].message.content




# 스테이블 디퓨전 기본 이미지 API 요청
import requests
import json
def sd_call(prompt, width, height):
    url = "https://stablediffusionapi.com/api/v3/text2img"
    payload = json.dumps({
        "key": "xbYWUXEs8U4iUweAQjOoH6mrasutoExqC3a2H6NtlnYAWTRXui9y0t2w1g8P",
        "prompt": prompt,
        "negative_prompt": "ng_deepnegative_v1_75t, (worst quality:1.4), (low quality:1.4), (normal quality:1.4), lowres, (nsfw:1.4)",
        "width": width,
        "height": height,
        "samples": "1",
        "num_inference_steps": "20",
        "seed": None,
        "guidance_scale": 7.5,
        "safety_checker": "yes",
        "multi_lingual": "no",
        "panorama": "no",
        "self_attention": "no",
        "upscale": "no",
        "embeddings_model": None,
        "webhook": None,
        "track_id": None
    })
    headers = {
        'Content-Type': 'application/json'
    }
    response = requests.request("POST", url, headers=headers, data=payload)
    response_data = json.loads(response.text)
    if response_data["status"] == "success":
        image_link = response_data["output"][0]
        fetch_result = "이미지 생성이 완료되었습니다"
    elif response_data["status"] == "processing":
        image_link = None
        fetch_result = response_data["fetch_result"]
    else:
        image_link = None
        fetch_result = "실패입니다 다시 실행해 주세요"
    return image_link, fetch_result



def sd_recall(fetch_result):
    url = fetch_result
    payload = json.dumps({
        "key": "xbYWUXEs8U4iUweAQjOoH6mrasutoExqC3a2H6NtlnYAWTRXui9y0t2w1g8P"
    })
    headers = {
        'Content-Type': 'application/json'
    }
    response = requests.request("POST", url, headers=headers, data=payload)
    response_data = json.loads(response.text)
    image_link = response_data["output"]

    if not image_link:
        image_link = None
    else:
        image_link = image_link[0]
    return image_link





with gr.Blocks(theme=gr.themes.Default()) as app:
    with gr.Tab("삽화 생성"):
        with gr.Row():
            # 1
            gr.Markdown(
            value="""
            # 삽화 생성
            짧은 소설에 어울리는 이미지를 생성할 수 있습니다.
            """
            )
        with gr.Row():
            # column1
            with gr.Column(scale=5):
                # 2
                pos_prompt = gr.Textbox(
                    label="이미지 생성 프롬프트를 작성해 주세요",
                    value="ultra realistic close up portrait ((beautiful pale cyberpunk female with heavy black eyeliner))",
                    lines=8,
                    interactive=True,
                )
                with gr.Row():
                    # 3
                    auto_prompt_generator = gr.Textbox(
                        label="자동 프롬프트 생성",
                        lines=6,
                        placeholder="이미지 생성을 위한 소설 내용을 작성해 주세요.\n자동으로 프롬프트가 생성됩니다.",
                        scale=7,
                    )
                    # 4
                    prompt_generator_btn = gr.Button(scale=1, value="자동\n생성")

                with gr.Group():
                    with gr.Row():
                        # 5
                        img_width = gr.Slider(
                            label="=Width", maximum=1024, value=512, interactive=True
                        )
                        # 6
                        img_height = gr.Slider(
                            label="=Height", maximum=1024, value=512, interactive=True
                        )

            # column2
            with gr.Column(scale=3):
                # 7
                output_status = gr.Textbox(
                    show_label=False, lines=1, placeholder="이미지 상태가 출력됩니다."
                )
                # 8
                generator_img = gr.Image(
                    value="https://pub-3626123a908346a7a8be8d9295f44e26.r2.dev/generations/8c595b57-563c-4417-9bfb-96aaebbb30b3-0.png",
                    label="이미지가 생성됩니다.",
                )
                # 9
                generator_img_btn = gr.Button(value="이미지 생성")
                # 10
                refresh_img_btn = gr.Button(value="이미지 새로고침")
        # 자동 생성 버튼 클릭
        prompt_generator_btn.click(
            fn=novel_keyword,
            inputs=[auto_prompt_generator],
            outputs=[pos_prompt]
        ) # 이미지 생성 버튼 클릭
        generator_img_btn.click(
            fn=sd_call,
            inputs=[pos_prompt, img_width, img_height],
            outputs=[generator_img, output_status]
        ) # 이미지 새로고침 버튼 클릭
        refresh_img_btn.click(
            fn=sd_recall,
            inputs=[output_status],
            outputs=[generator_img]
        )
    with gr.Tab("이미지 편집") as edit_tab:
        with gr.Row():
            #1
            gr.Markdown(
            value="""
            # 이미지 편집
            생성한 이미지를 편집할 수 있습니다.
            """)
        with gr.Row():
            #2
            edit_prompt = gr.Textbox(
                label="이미지 수정 프롬프트를 작성해 주세요",
                value="black hair",
                lines=5,
                interactive=True,
                scale=7
            )
            #3
            edit_btn = gr.Button(
                value="이미지 편집",
                scale=1
            )
        with gr.Row():
            #4
            edit_status = gr.Textbox(
                show_label=False,
                lines=1,
                placeholder="편집 이미지 상태가 출력됩니다.",
                scale=7
            )
            #5
            refresh_edit_btn = gr.Button(
                value="이미지 새로고침",
                scale=1
            )
        with gr.Row():
            #6
            org_img = gr.ImageMask(
                label="수정 전 이미지",
                image_mode='RGB',
                brush = gr.Brush(
                    default_size=20,
                    colors=["#FFFFFF"],
                    color_mode="fixed",
                    ),
                show_label=True
                )
            #7
            edit_img = gr.Image(
                label="수정 후 이미지",
            )
    edit_tab.select(
        fn=edit_load_img,
        inputs=[generator_img],
        outputs=[org_img]
    )
    #이미지 편집 버튼 클릭
    edit_btn.click(
        fn=edit_img_generator,
        inputs=[org_img, edit_prompt],
        outputs=[edit_img, edit_status]
    )
    #이미지
    refresh_edit_btn.click(
        fn=sd_edit_recall,
        inputs=[edit_status],
        outputs=[edit_img]
    )
app.launch(debug=True)