import os
import gradio as gr
import requests

# 从新版 openai>=1.0.0 库导入
from openai import OpenAI

##############################################################################
# 1. Furry 物种数据
##############################################################################
furry_species_map = {
    "CANIDS (Canines)": ["Dogs", "Wolves", "Foxes", "Jackals"],
    "FELINES (Cats)": ["HouseCats", "Lions", "Tigers", "Cheetahs"],
    "CUSTOM (Just An Example)": ["Dragons", "Werewolves", "Kitsune"],
    # ... 这里你可以替换成更完整的物种分类，或保留最简
}

def flatten_species(map_dict):
    """
    将分门别类的物种字典展开成一个扁平列表，比如:
    ["CANIDS (Canines) - Dogs", "CANIDS (Canines) - Wolves", ...]
    """
    result = []
    for category, species_list in map_dict.items():
        for sp in species_list:
            result.append(f"{category} - {sp}")
    return sorted(result)

ALL_FURRY_SPECIES = flatten_species(furry_species_map)

##############################################################################
# 2. 核心：调用 GPT 或 DeepSeek 生成描述 & 翻译
##############################################################################
def generate_transformed_prompt(prompt, gender_option, furry_species, api_mode, api_key):
    """
    1) 构造 tags（包括性别/物种）
    2) 选择 GPT / DeepSeek 调用
    3) 返回生成的描述
    """
    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

    # 2. 根据选择调用 GPT / DeepSeek
    if api_mode == "GPT":
        # GPT
        base_url = None
        model_name = "gpt-3.5-turbo"  # 你可改成 "gpt-4" 等
    else:
        # DeepSeek
        base_url = "https://api.deepseek.com"
        model_name = "deepseek-chat"

    # 创建客户端
    if not api_key:
        return "Error: API Key not provided."

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

    # 拼出文字供对话
    tag_desc = "\n".join([f"{k}: {v}" for k, v in tags.items() if v])

    try:
        # 发起 chat.completions.create()
        response = client.chat.completions.create(
            model=model_name,
            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:\n{tag_desc}\nPlease generate a vivid, imaginative scene description.",
                },
            ]
        )
        return response.choices[0].message.content.strip()
    except Exception as e:
        return f"{api_mode} generation failed. Error: {e}"


def translate_prompt_text(text, target_language, api_mode, api_key):
    """
    调用 GPT / DeepSeek 做翻译，简化写法：同样只需改 model/base_url
    """
    if not api_key:
        return "Error: API Key not provided."

    if not text.strip():
        return ""

    if api_mode == "GPT":
        base_url = None
        model_name = "gpt-3.5-turbo"
    else:
        base_url = "https://api.deepseek.com"
        model_name = "deepseek-chat"

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

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

##############################################################################
# 3. Gradio 界面
##############################################################################
def build_app():
    with gr.Blocks() as demo:
        gr.Markdown("## Prompt TransTool - GPT/DeepSeek 性别物种转换")

        with gr.Row():
            with gr.Column():
                api_mode = gr.Radio(
                    label="选择 API（GPT / DeepSeek）",
                    choices=["GPT", "DeepSeek"],
                    value="GPT",
                )
                api_key = gr.Textbox(
                    label="API 密钥 (API Key)",
                    type="password",
                    placeholder="请输入 GPT 或 DeepSeek 的 API 密钥"
                )
                gender_option = gr.Radio(
                    label="转换目标 (Gender / Furry)",
                    choices=[
                        "Trans_to_Male",
                        "Trans_to_Female",
                        "Trans_to_Mannequin",
                        "Trans_to_Intersex",
                        "Trans_to_Furry",
                    ],
                    value="Trans_to_Male",
                )
                furry_species_box = gr.Dropdown(
                    label="Furry 物种 (Furry Species)",
                    choices=ALL_FURRY_SPECIES,
                    value=None,
                    visible=False
                )
                def show_furry_spec(g):
                    return gr.update(visible=(g == "Trans_to_Furry"))
                gender_option.change(show_furry_spec, inputs=[gender_option], outputs=[furry_species_box])

            with gr.Column():
                prompt_input = gr.Textbox(
                    label="提示词 (Prompt)",
                    lines=5,
                    placeholder="示例：一位穿红色连衣裙的少女在花园中..."
                )
                gen_output = gr.Textbox(
                    label="转换后 (Transformed Prompt)",
                    lines=6
                )

        # 翻译
        with gr.Row():
            translate_language = gr.Dropdown(
                label="翻译语言 (Translation Language)",
                choices=["English", "Chinese", "Japanese", "French", "German", "Spanish"],
                value="English",
            )
            translate_output = gr.Textbox(
                label="翻译结果 (Translation Result)",
                lines=6
            )

        def on_generate(prompt, gender, furry, mode, key, lang):
            # 1) 生成转换后提示词
            transformed = generate_transformed_prompt(prompt, gender, furry, mode, key)
            # 2) 翻译
            translated = translate_prompt_text(transformed, lang, mode, key)
            return transformed, translated

        prompt_input.submit(
            fn=on_generate,
            inputs=[prompt_input, gender_option, furry_species_box, api_mode, api_key, translate_language],
            outputs=[gen_output, translate_output],
        )
        btn = gr.Button("生成 / Generate")
        btn.click(
            fn=on_generate,
            inputs=[prompt_input, gender_option, furry_species_box, api_mode, api_key, translate_language],
            outputs=[gen_output, translate_output],
        )

    return demo

if __name__ == "__main__":
    demo = build_app()
    demo.launch()