import gradio as gr
import openai
import json
import os
import random

# OpenAI API 클라이언트 설정
openai.api_key = os.getenv("OPENAI_API_KEY")

# JSON 데이터 로드
json_path = "back.json"  # JSON 파일 경로
with open(json_path, "r", encoding="utf-8") as f:
    data = json.load(f)

# JSON 데이터를 기반으로 카테고리와 배경 목록 생성
categories = {}
for item in data:
    category = item.get("카테고리", "Unknown")
    korean_translation = item.get("Korean Translation", "Unknown")
    english_prompt = item.get("English Prompt", "Unknown")

    if category not in categories:
        categories[category] = []
    categories[category].append({"korean": korean_translation, "english": english_prompt})

# LLM 호출 함수
def call_api(content, system_message, max_tokens, temperature, top_p):
    response = openai.ChatCompletion.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": system_message},
            {"role": "user", "content": content},
        ],
        max_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
    )
    return response.choices[0].message['content'].strip()

# 사물 및 배경 번역 함수
def translate_object(object_input):
    system_message = "Translate the following object description from Korean to English. Keep it simple and concise."
    return call_api(
        content=object_input,
        system_message=system_message,
        max_tokens=10,
        temperature=0.7,
        top_p=1.0
    )

def translate_background(background_input):
    system_message = "Translate the following background description from Korean to English. Keep it simple and concise."
    return call_api(
        content=background_input,
        system_message=system_message,
        max_tokens=20,
        temperature=0.7,
        top_p=1.0
    )

# 직접 배경 입력 탭의 프롬프트 생성 함수
def generate_prompt(object_input, background_input):
    translated_object = translate_object(object_input)
    translated_background = translate_background(background_input)
    return f"advertisement photography, {translated_object} {translated_background}"

# 1차/2차 카테고리 선택 탭의 프롬프트 생성 함수
def generate_category_prompt(object_input, category, background):
    translated_object = translate_object(object_input)
    background_english = next(
        (item["english"] for item in categories.get(category, []) if item["korean"] == background),
        ""
    )
    if translated_object and background_english:
        return f"advertisement photography, {translated_object} {background_english}."
    else:
        return "Error: Please ensure both object and background are provided correctly."

# Gradio 인터페이스 생성
with gr.Blocks() as demo:
    with gr.Tab("직접 배경 입력"):
        gr.Markdown("### 직접 배경 입력을 통한 프롬프트 생성")

        with gr.Row():
            object_input = gr.Textbox(label="사물 입력", placeholder="예: 오렌지 주스")
            background_input = gr.Textbox(label="배경 입력", placeholder="예: 깨끗한 주방")
            output = gr.Textbox(label="생성된 영어 프롬프트")

        generate_button = gr.Button("프롬프트 생성")
        generate_button.click(
            fn=generate_prompt,
            inputs=[object_input, background_input],
            outputs=output
        )
    
    with gr.Tab("1차/2차 카테고리 선택"):
        gr.Markdown("### 1차/2차 카테고리 선택을 통한 프롬프트 생성")

        with gr.Row():
            object_input = gr.Textbox(label="사물 입력", placeholder="예: 오렌지 주스")
            category_dropdown = gr.Dropdown(
                label="1차 카테고리 선택",
                choices=list(categories.keys()),
                interactive=True
            )
            background_dropdown = gr.Dropdown(
                label="2차 배경 선택",
                choices=[],
                interactive=True
            )

            def update_background_options(category):
                if category in categories:
                    return gr.update(choices=[item["korean"] for item in categories[category]], value=None)
                else:
                    return gr.update(choices=[], value=None)

            category_dropdown.change(
                fn=update_background_options,
                inputs=category_dropdown,
                outputs=background_dropdown
            )
        
        generate_button = gr.Button("프롬프트 생성")
        category_prompt_output = gr.Textbox(label="생성된 영어 프롬프트")

        generate_button.click(
            fn=generate_category_prompt,
            inputs=[object_input, category_dropdown, background_dropdown],
            outputs=category_prompt_output
        )

demo.launch()