##############################################
# app.py
##############################################
import os
import requests
import gradio as gr

# 新版 openai>=1.0.0 引用方式
from openai import OpenAI

##############################################################################
# 1. GPT 调用示例函数（新版 openai>=1.0.0）
##############################################################################
def generate_natural_language_description_gpt(tags, api_key=None, base_url=None, model="gpt-4o"):
    """
    使用新版 openai>=1.0.0 库来调用 GPT。
      1) from openai import OpenAI
      2) client = OpenAI(api_key=...)
      3) client.chat.completions.create(...)
    """
    # 若没传入 api_key，就从环境变量获取
    if not api_key:
        api_key = os.getenv("OPENAI_API_KEY")
    if not api_key:
        return "Error: No GPT API Key provided."

    # 创建一个新的 OpenAI Client
    client = OpenAI(api_key=api_key)
    
    # 如果你有自定义 base_url（如代理/私有化部署），可在此设置
    if base_url:
        client.base_url = base_url

    # 将 tags 拼出可读字符串
    tag_descriptions = "\n".join([
        f"{k}: {', '.join(v) if isinstance(v, list) else v}"
        for k, v in tags.items() if v
    ])

    try:
        chat_completion = client.chat.completions.create(
            model=model,  # 例如 "gpt-4o", "gpt-4o-mini", "gpt-3.5-turbo" 等
            messages=[
                {
                    "role": "system",
                    "content": (
                        "You are a creative assistant that generates detailed and imaginative scene descriptions "
                        "for AI generation prompts. Focus on the details provided and incorporate them into a "
                        "cohesive narrative. Use at least three sentences but no more than five sentences."
                    ),
                },
                {
                    "role": "user",
                    "content": f"Here are the tags and details:\n{tag_descriptions}\nPlease generate a vivid, imaginative scene description.",
                },
            ]
        )
        # 新版库返回对象中，内容在 chat_completion.choices[0].message.content
        return chat_completion.choices[0].message.content.strip()
    except Exception as e:
        return f"GPT generation failed. Error: {e}"

##############################################################################
# 2. DeepSeek 调用示例函数（真实 HTTP 请求示例，需你修改）
##############################################################################
def generate_natural_language_description_deepseek(tags, api_key=None, base_url=None):
    """
    调用 DeepSeek API 获取描述。如果你不需要 DeepSeek，可删除此函数。
    注意：此处示例的接口 URL、请求体和返回格式需要替换为真实 DeepSeek 文档内容。
    """
    if not api_key:
        return "Error: No DeepSeek API Key provided."

    # 假设你的 DeepSeek 接口地址如下，请改为实际地址
    url = base_url or "https://api.deepseek.com/v1/generate"

    try:
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {api_key}"
        }
        payload = {
            "tags": tags
        }
        resp = requests.post(url, json=payload, headers=headers, timeout=30)
        if resp.status_code == 200:
            data = resp.json()
            # 根据你们的返回结构来取值
            if data.get("success"):
                return data["data"].get("description", "No description found.")
            else:
                return f"DeepSeek generation success=false. {data}"
        else:
            return f"DeepSeek generation failed. Status {resp.status_code}, {resp.text}"
    except Exception as e:
        return f"DeepSeek generation request error: {e}"

##############################################################################
# 3. GPT 翻译函数（新版 openai>=1.0.0）
##############################################################################
def translate_text_with_gpt(text, target_language, api_key=None, base_url=None, model="gpt-4o"):
    if not api_key:
        api_key = os.getenv("OPENAI_API_KEY")
    if not api_key:
        return "Error: No GPT Translation Key provided."

    client = OpenAI(api_key=api_key)
    if base_url:
        client.base_url = base_url

    try:
        # Prompt: 让 GPT 扮演翻译
        system_prompt = f"You are a professional translator. Translate the following text to {target_language}:"
        chat_completion = client.chat.completions.create(
            model=model,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user",   "content": text},
            ]
        )
        return chat_completion.choices[0].message.content.strip()
    except Exception as e:
        return f"GPT translation failed. Error: {e}"

##############################################################################
# 4. DeepSeek 翻译函数（需你修改为真实接口）
##############################################################################
def translate_text_with_deepseek(text, target_language, api_key=None, base_url=None):
    if not api_key:
        return "Error: No DeepSeek Translation Key provided."

    # 假设 DeepSeek 翻译接口如下，请改为实际
    url = base_url or "https://api.deepseek.com/v1/translate"

    try:
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {api_key}"
        }
        payload = {
            "text": text,
            "target_language": target_language
        }
        resp = requests.post(url, json=payload, headers=headers, timeout=30)
        if resp.status_code == 200:
            data = resp.json()
            if data.get("success"):
                return data["data"].get("translated_text", "No translated_text found.")
            else:
                return f"DeepSeek translation success=false. {data}"
        else:
            return f"DeepSeek translation failed. Status {resp.status_code}, {resp.text}"
    except Exception as e:
        return f"DeepSeek translation request error: {e}"

##############################################################################
# 5. 根据用户选择进行提示词转换并调用 GPT/DeepSeek 生成描述
##############################################################################
def transform_prompt(prompt, gender_option, furry_species, api_mode, api_key):
    """
    示例转换逻辑：将 prompt 结合 gender/furry 信息组成 tags，然后调用指定 API。
    你可以在此处扩充更复杂的处理逻辑。
    """
    tags = {}
    
    # 性别 / 物种 设定
    if gender_option == "Trans_to_Male":
        tags["gender"] = "male"
    elif gender_option == "Trans_to_Female":
        tags["gender"] = "female"
    elif gender_option == "Trans_to_Mannequin":
        tags["gender"] = "genderless"
    elif gender_option == "Trans_to_Intersex":
        tags["gender"] = "intersex"
    elif gender_option == "Trans_to_Furry":
        tags["gender"] = "furry"
        tags["furry_species"] = furry_species or "unknown"
    
    # 原始 prompt
    tags["base_prompt"] = prompt

    # 调用对应 API
    if api_mode == "GPT":
        scene_description = generate_natural_language_description_gpt(tags, api_key)
    else:  # DeepSeek
        scene_description = generate_natural_language_description_deepseek(tags, api_key)
    
    return scene_description

##############################################################################
# 6. 翻译逻辑：看用户选择，用 GPT 或 DeepSeek 做翻译
##############################################################################
def do_translation(scene_desc, translate_language, api_mode, api_key):
    if not scene_desc.strip():
        return ""
    if api_mode == "GPT":
        return translate_text_with_gpt(scene_desc, translate_language, api_key)
    else:
        return translate_text_with_deepseek(scene_desc, translate_language, api_key)

##############################################################################
# 7. 搭建 Gradio 界面
##############################################################################
def build_interface():
    with gr.Blocks() as demo:

        gr.Markdown("## Prompts_TransTool - 提示词一键性别物种转换器（openai>=1.0.0 版）")

        with gr.Row():
            with gr.Column():
                # 选择 API 服务（GPT / DeepSeek）
                api_mode = gr.Radio(
                    label="选择 API 服务 (Choose API Service)",
                    choices=["GPT", "DeepSeek"],
                    value="GPT"
                )

                # 输入 API Key
                api_key = gr.Textbox(
                    label="API 密钥 (API Key)",
                    type="password",
                    placeholder="在此输入 GPT 或 DeepSeek 的 API 密钥"
                )

                # 性别 / Furry 选项
                gender_option = gr.Radio(
                    label="性别 / Furry 选项",
                    choices=[
                        "Trans_to_Male",
                        "Trans_to_Female",
                        "Trans_to_Mannequin",
                        "Trans_to_Intersex",
                        "Trans_to_Furry",
                    ],
                    value="Trans_to_Male",
                )

                # 如果是 Furry，显示可选物种
                furry_species = gr.Dropdown(
                    label="Furry 物种 (Furry Species)",
                    choices=["Wolf", "Fox", "Tiger", "Lion"],
                    value=None,
                    visible=False
                )

                def show_furry_species(gender):
                    return gr.update(visible=(gender == "Trans_to_Furry"))
                gender_option.change(
                    fn=show_furry_species,
                    inputs=[gender_option],
                    outputs=[furry_species]
                )

            with gr.Column():
                # 用户输入 prompt
                user_prompt = gr.Textbox(
                    label="提示词 (Prompt)",
                    lines=5,
                    placeholder=(
                        "示例：一位穿着红色连衣裙的少女，坐在落日余晖下的草地上..."
                    )
                )

                # 转换后输出
                generated_output = gr.Textbox(
                    label="转换后的提示词 (Generated Trans-Description)",
                    lines=7
                )

        with gr.Row():
            # 翻译语言
            translate_language = gr.Dropdown(
                label="翻译语言 (Translation Language)",
                choices=[
                    "English", "Chinese", "Japanese", "French", "German", 
                    "Dutch", "Arabic", "Russian", "Persian", "Italian"
                ],
                value="English",
            )
            # 翻译结果
            translated_text = gr.Textbox(
                label="翻译结果 (Translated Result)",
                lines=7
            )

        ############################################################################
        # 事件逻辑
        ############################################################################
        # 1) 生成并翻译
        def on_generate(prompt, gender, furry, mode, key, lang):
            trans_desc = transform_prompt(prompt, gender, furry, mode, key)
            trans_result = do_translation(trans_desc, lang, mode, key)
            return trans_desc, trans_result

        # 回车提交
        user_prompt.submit(
            fn=on_generate,
            inputs=[user_prompt, gender_option, furry_species, api_mode, api_key, translate_language],
            outputs=[generated_output, translated_text],
        )

        # 按钮点击
        generate_button = gr.Button("生成 / Generate")
        generate_button.click(
            fn=on_generate,
            inputs=[user_prompt, gender_option, furry_species, api_mode, api_key, translate_language],
            outputs=[generated_output, translated_text],
        )

        # 2) 语言切换再翻译
        def on_translate(scene_desc, lang, mode, key):
            return do_translation(scene_desc, lang, mode, key)

        translate_language.change(
            fn=on_translate,
            inputs=[generated_output, translate_language, api_mode, api_key],
            outputs=[translated_text]
        )

    return demo

##############################################################################
# 主入口
##############################################################################
if __name__ == "__main__":
    demo = build_interface()
    demo.launch()