import os 
import gradio as gr
import openai

##############################################################################
# 1. GPT 或 DeepSeek 调用示例函数
##############################################################################
def generate_natural_language_description_gpt(tags, api_key, base_url=None, model="gpt-4"):
    """
    使用 OpenAI GPT 生成自然语言描述的示例函数。
    """
    if not api_key:
        return "Error: GPT API Key not provided."

    # 设置 API
    openai.api_key = api_key
    if base_url:
        openai.api_base = base_url

    # 将 dict 转成可读字符串
    tag_descriptions = "\n".join([
        f"{key}: {', '.join(value) if isinstance(value, list) else value}"
        for key, value in tags.items() if value
    ])

    try:
        response = openai.ChatCompletion.create(
            model=model,
            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.",
                },
            ]
        )
        return response.choices[0].message.content.strip()
    except Exception as e:
        return f"GPT generation failed. Error: {e}"


def generate_natural_language_description_deepseek(tags, api_key, base_url=None):
    """
    使用 DeepSeek API 生成自然语言描述的示例函数。
    这里演示伪代码，你需要根据实际 DeepSeek 的文档进行实现。
    """
    if not api_key:
        return "Error: DeepSeek API Key not provided."
    
    # 伪代码示例（需根据你的 DeepSeek API 文档做实际实现）
    # import requests
    # response = requests.post(
    #     url=base_url or "https://api.deepseek.com/xxx", 
    #     headers={"Authorization": f"Bearer {api_key}"},
    #     json={"tags": tags}
    # )
    # return response.json()["description"]
    return "DeepSeek 生成的描述（此处为示例伪代码）"


##############################################################################
# 2. 翻译示例函数（使用 GPT 或 DeepSeek）
##############################################################################
def translate_text_with_gpt(text, target_language, api_key, base_url=None, model="gpt-4"):
    """
    使用 GPT 来进行翻译的简单示例。
    """
    if not api_key:
        return "Error: GPT Translation Key not provided."

    openai.api_key = api_key
    if base_url:
        openai.api_base = base_url

    try:
        system_prompt = f"You are a professional translator. Translate the following text to {target_language}:"
        response = openai.ChatCompletion.create(
            model=model,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": text},
            ]
        )
        return response.choices[0].message.content.strip()
    except Exception as e:
        return f"GPT translation failed. Error: {e}"


def translate_text_with_deepseek(text, target_language, api_key, base_url=None):
    """
    使用 DeepSeek 来进行翻译的简单示例（伪代码）。
    """
    if not api_key:
        return "Error: DeepSeek Translation Key not provided."
    # 同样需要根据 DeepSeek 的文档来实现
    return f"DeepSeek翻译后的文本（示例）。目标语言：{target_language}"


##############################################################################
# 3. 根据用户选择进行提示词转换并调用 GPT/DeepSeek 生成描述
##############################################################################
def transform_prompt(prompt, gender_option, furry_species, api_mode, api_key):
    """
    性别/物种转换的简单示例逻辑，然后调用相应 API。
    你可在此处结合“关于Male/Female/ambiguous/intersex的details”添加更复杂的处理。
    """
    tags = {}
    
    # 根据选择设置性别或物种标签
    if gender_option == "Trans_to_Male":
        # 这里可以参考你的细节 rules 做更加复杂的转换
        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"
    
    # 原始提示词
    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


##############################################################################
# 4. 调用翻译函数
##############################################################################
def do_translation(scene_desc, translate_language, api_mode, api_key):
    """
    根据选择的 API（GPT/DeepSeek）进行翻译。
    """
    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)


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

        gr.Markdown("## Prompts_TransTool-提示词一键性别物种转换器")

        with gr.Row():
            with gr.Column():
                # 选择调用哪个 API
                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 选项 (Gender / 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  # 初始不可见
                )

                # 当性别选项切换时，如果选择 Furry，就显示物种下拉，否则隐藏
                def show_furry_species(gender):
                    return gr.update(visible=(gender == "Trans_to_Furry"))

                gender_option.change(
                    show_furry_species, 
                    inputs=[gender_option], 
                    outputs=[furry_species]
                )

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

                # 输出场景描述
                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
            )

        ######################################################################
        # 事件绑定
        ######################################################################
        
        # 新增：生成时，直接返回「转换结果」和「翻译结果」并一起更新
        def on_generate(prompt, gender, furry, mode, key, lang):
            # 1) 先做性别/物种转换，拿到“转换后”提示词
            trans_desc = transform_prompt(prompt, gender, furry, mode, key)
            # 2) 立刻翻译
            trans_result = do_translation(trans_desc, lang, mode, key)
            # 返回两项
            return trans_desc, trans_result

        # 当用户在 prompt 输入后按回车时，触发生成场景描述 + 翻译
        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],
        )

        # 当用户切换翻译语言时，如果已经有转换后的内容，则再翻译一次
        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


# 在 Spaces 启动
if __name__ == "__main__":
    demo = build_interface()
    demo.launch()