import os 
import gradio as gr
import openai
import requests

##############################################################################
# 1. GPT 调用示例函数
##############################################################################
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."

    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}"

##############################################################################
# 2. DeepSeek 调用示例函数 (真实请求)
##############################################################################
def generate_natural_language_description_deepseek(tags, api_key, base_url=None):
    """
    使用 DeepSeek API 生成自然语言描述的真实示例函数。
    假设 DeepSeek 文档里的生成接口是:
      POST https://api.deepseek.com/v1/generate
      Headers: {"Authorization": "Bearer <api_key>"}
      Body(JSON): {"tags": {...}}
    返回: 
      {"success": true, "data": {"description": "..."}}
    """
    if not api_key:
        return "Error: DeepSeek API Key not provided."
    
    # 如果有自定义 base_url 就用，没有就用假设的默认
    url = base_url or "https://api.deepseek.com/v1/generate"

    try:
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {api_key}"
        }
        payload = {
            "tags": tags  # 具体字段要根据实际 DeepSeek 文档
        }
        resp = requests.post(url, json=payload, headers=headers, timeout=30)
        if resp.status_code == 200:
            j = resp.json()
            if j.get("success"):
                # 假设描述存放在 data.description
                return j["data"].get("description", "No 'description' found in data.")
            else:
                return f"DeepSeek generation failed, success=false. {j}"
        else:
            return f"DeepSeek generation failed with status {resp.status_code}. {resp.text}"
    except Exception as e:
        return f"DeepSeek generation request error: {e}"

##############################################################################
# 3. GPT 翻译函数
##############################################################################
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}"

##############################################################################
# 4. DeepSeek 翻译函数 (真实请求)
##############################################################################
def translate_text_with_deepseek(text, target_language, api_key, base_url=None):
    """
    使用 DeepSeek 来进行翻译的真实示例函数。
    假设 DeepSeek 文档里的翻译接口是:
      POST https://api.deepseek.com/v1/translate
      Headers: {"Authorization": "Bearer <api_key>"}
      Body(JSON): {"text": "...", "target_language": "..."}
    返回: 
      {"success": true, "data": {"translated_text": "..."}}
    """
    if not api_key:
        return "Error: DeepSeek Translation Key not provided."
    
    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:
            j = resp.json()
            if j.get("success"):
                # 假设翻译结果存放在 data.translated_text
                return j["data"].get("translated_text", "No 'translated_text' found in data.")
            else:
                return f"DeepSeek translation failed, success=false. {j}"
        else:
            return f"DeepSeek translation failed with 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):
    """
    性别/物种转换的简单示例逻辑，然后调用相应 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"
    
    # 原始提示词
    tags["base_prompt"] = prompt

    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. 调用翻译函数
##############################################################################
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)

##############################################################################
# 7. 搭建 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) 调用 transform_prompt 拿到转换后的描述
            trans_desc = transform_prompt(prompt, gender, furry, mode, key)
            # 2) 调用翻译
            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],
        )

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