app.py

import os
import json
import re
import gradio as gr
from openai import OpenAI

##############################################################################
# 1. 读取外部文件
##############################################################################
try:
    with open("furry_species.json", "r", encoding="utf-8") as ff:
        FURRY_DATA = json.load(ff)
except:
    FURRY_DATA = {}

try:
    with open("gender_rules.json", "r", encoding="utf-8") as gf:
        GENDER_RULES = json.load(gf)
except:
    GENDER_RULES = {}

try:
    with open("transform_rules.json", "r", encoding="utf-8") as tf:
        TRANSFORM_DICT = json.load(tf)
except:
    TRANSFORM_DICT = {}

##############################################################################
# 2. 多级菜单函数
##############################################################################
def get_top_categories(furry_data):
    return sorted(list(furry_data.keys()))

def get_sub_categories(furry_data, top_category):
    if top_category in furry_data:
        return sorted(list(furry_data[top_category].keys()))
    return []

def get_species_list(furry_data, top_category, sub_category):
    if top_category in furry_data and sub_category in furry_data[top_category]:
        species = furry_data[top_category][sub_category]
        # 检查 species 是否是列表，并且每个元素是否有 "Name" 字段
        if isinstance(species, list) and all(isinstance(item, dict) and "Name" in item for item in species):
            return [item ["Name"] for item in species]
        else:
            print(f"[DEBUG] Unexpected structure for species: {species}")
    return []

##############################################################################
# 3. 合并规则文本
##############################################################################
def merge_transform_rules_into_prompt(rules_json):
    if not rules_json:
        return "(No transform rules loaded)"
    
    gt = rules_json.get("gender_transform", {})
    sp = rules_json.get("shared_preferences", {})
    td = rules_json.get("table_details", {})

    text_parts = []
    text_parts.append("==== GENDER TRANSFORM RULES ====")
    text_parts.append(str(gt))

    text_parts.append("==== SHARED PREFERENCES ====")
    text_parts.append(str(sp))

    text_parts.append("==== TABLE DETAILS (PRO ACTIONS) ====")
    text_parts.append(str(td))

    return "\n".join(text_parts)

RULES_TEXT_FULL = merge_transform_rules_into_prompt(TRANSFORM_DICT)

##############################################################################
# 4. 强制替换：根据 override_conflicting_descriptors
##############################################################################

transform_map = {
    "Trans_to_Male": "female_to_male",
    "Trans_to_Female": "male_to_female",
    "Trans_to_Mannequin": "any_to_genderless",
    "Trans_to_Intersex": "any_to_intersex",
    "Trans_to_Furry": "trans_to_furry"
}

def forced_replace(prompt, direction):
    if not TRANSFORM_DICT:
        return prompt

    override_section = TRANSFORM_DICT.get("override_conflicting_descriptors", {})
    replacements = override_section.get(direction, {})
    if not replacements:
        return prompt

    for old, new in replacements.items():
        pattern = r"(?i)\b" + re.escape(old) + r"\b"
        prompt = re.sub(pattern, new, prompt)

    # 针对复杂句子的补充替换
    if direction == "Trans_to_male":
        # 女性到男性
        prompt = re.sub(r"\bshe\b", "he", prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"\bher\b", "his", prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"\bherself\b", "himself", prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"\bwomen\b", "men", prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"\blady\b", "gentleman", prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"\bqueen\b", "king", prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"\bfemale\b", "male", prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"\bgirl\b", "boy", prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"\bprincess\b", "prince", prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"\bactress\b", "actor", prompt, flags=re.IGNORECASE)
    
    elif direction == "Trans_to_female":
        # 男性到女性
        prompt = re.sub(r"\bhe\b", "she", prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"\bhis\b", "her", prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"\bhimself\b", "herself", prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"\bmen\b", "women", prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"\bgentleman\b", "lady", prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"\bking\b", "queen", prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"\bmale\b", "female", prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"\bboy\b", "girl", prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"\bprince\b", "princess", prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"\bactor\b", "actress", prompt, flags=re.IGNORECASE)
    
    elif direction == "Trans_to_mannequin":
        # 转换为无性别或人偶化
        prompt = re.sub(r"\bshe\b|\bhe\b", "it", prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"\bher\b|\bhis\b", "its", prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"\bherself\b|\bhimself\b", "itself", prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"\bwoman\b|\bman\b", "mannequin-like figure", prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"\bgirl\b|\bboy\b", "character", prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"\bqueen\b|\bking\b", "figurehead", prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"\bfemale\b|\bmale\b", "genderless", prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"\bprincess\b|\bprince\b", "androgynous heir", prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"\bactress\b|\bactor\b", "performer", prompt, flags=re.IGNORECASE)
    
    elif direction == "Trans_to_intersex":
        # 转换为双性化特征
        prompt = re.sub(r"\bshe\b", "they", prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"\bhe\b", "they", prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"\bher\b", "their", prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"\bhis\b", "their", prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"\bherself\b", "themself", prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"\bhimself\b", "themself", prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"\bwoman\b", "androgynous individual", prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"\bman\b", "androgynous individual", prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"\bgirl\b", "intersex character", prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"\bboy\b", "intersex character", prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"\bfemale\b", "intersex", prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"\bmale\b", "intersex", prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"\bqueen\b", "intersex ruler", prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"\bking\b", "intersex ruler", prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"\bprincess\b", "intersex heir", prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"\bprince\b", "intersex heir", prompt, flags=re.IGNORECASE)

    return prompt

##############################################################################
# 5. 核心 GPT/DeepSeek 调用
##############################################################################
def generate_transformed_output(prompt, gender_option, top_cat, sub_cat, species_item, api_mode, api_key):
    if not api_key:
        return "Error: No API Key provided."

    if api_mode == "GPT":
        base_url = None
        model_name = "gpt-4o"
    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

    if gender_option == "Trans_to_Furry":
        furry_path = f"{top_cat} > {sub_cat} > {species_item}" if (top_cat and sub_cat and species_item) else "unknown"
        species_data = FURRY_DATA.get(top_cat, {}).get(sub_cat, [])
        species_details = next((item["Details"] for item in species_data if item["Name"] == species_item), "unknown details")
        gender_desc = ", ".join([g.lower() for g in gender_selection_furry]) if gender_selection_furry else "neutral"
        extra_line = f"\nfurry {gender_desc} {furry_path} {species_details} anthro character\n"
    else:
        extra_line = ""

    gender_specific_rule = ""
    if gender_option == "Trans_to_Male":
        gender_specific_rule = GENDER_RULES.get("male", "")
        extra_line = f"\n1boy\n"
    elif gender_option == "Trans_to_Female":
        gender_specific_rule = GENDER_RULES.get("female", "")
        extra_line = f"\n1girl\n"
    elif gender_option == "Trans_to_Mannequin":
        gender_specific_rule = GENDER_RULES.get("genderless", "")
        extra_line = f"\nmannequin or statue or doll or robot or android\n"
    elif gender_option == "Trans_to_Intersex":
        gender_specific_rule = GENDER_RULES.get("intersex", "")
        extra_line = f"\nhanve both penis_testicles and pussy\n"

    system_prompt = f"""
You are a creative assistant that transforms the user's base prompt 
生成将提示词中的主体转化为{extra_line}{gender_specific_rule}的提示词，注意性别，其服装和装饰细节，皮肤和第二性征的转变，并且 Follow references below:
Detailed Transform Rules (transform_rules.json): 简化并去除重复和冗余的提示词部分，删除和主体角色性别或种族不符合的，不融洽的提示词。
{RULES_TEXT_FULL}
Instructions:
- Ensure all gender-specific terms are replaced correctly based on the rules.
- Original prompt tags: {prompt}
- Convert them into NEW combined tags, removing or replacing conflicting ones.
- Only output two parts:
  1) One line of final tags in parentheses, e.g. (male, short hair, dynamic pose, ...)
  2) A blank line.
  3) Then 2~5 sentences of imaginative scene description in English.
- No extra lines, no headings, no 'gender:' or 'base_prompt:'.
- No more than 990 tokens.
- End of instructions.
""".strip()

    try:
        resp = client.chat.completions.create(
            model=model_name,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user",   "content": "Generate final tags and description now."}
            ],
        )
        return resp.choices[0].message.content.strip()
    except Exception as e:
        return f"{api_mode} generation failed. Error: {e}"

##############################################################################
# 6. 翻译函数
##############################################################################
def translate_text(content, lang, api_mode, api_key):
    if not api_key:
        return "Error: No API Key provided."
    if not content.strip():
        return ""

    if api_mode == "GPT":
        base_url = None
        model_name = "gpt-4o"
    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

    translate_system_prompt = f"""
You are a translator. Translate the following text to {lang}, 
keeping parentheses line, tags and blank line if present. 
No extra headings.
""".strip()

    try:
        resp = client.chat.completions.create(
            model=model_name,
            messages=[
                {"role": "system", "content": translate_system_prompt},
                {"role": "user",   "content": content}
            ],
        )
        return resp.choices[0].message.content.strip()
    except Exception as e:
        return f"{api_mode} translation failed. Error: {e}"

##############################################################################
# 7. Gradio 界面
##############################################################################
def build_interface():
    with gr.Blocks() as demo:
        gr.Markdown("## Prompt Trans-Tool - 提示词物种性别转换器")

        with gr.Row():
            with gr.Column():
                api_mode = gr.Radio(
                    label="Select API 选择API厂商 (GPT/DeepSeek)",
                    choices=["GPT", "DeepSeek"],
                    value="GPT"
                )
                api_key = gr.Textbox(
                    label="API Key",
                    type="password",
                    placeholder="Input your GPT or DeepSeek Key"
                )

                gender_option = gr.Radio(
                    label="Trans-Option 选择转换目标",
                    choices=[
                        "Trans_to_Male",
                        "Trans_to_Female",
                        "Trans_to_Mannequin",
                        "Trans_to_Intersex",
                        "Trans_to_Furry"
                    ],
                    value="Trans_to_Male"
                )

                top_cat_dd = gr.Dropdown(
                    label="Furry: Top Category",
                    choices=get_top_categories(FURRY_DATA),
                    value=None,
                    visible=False
                )
                sub_cat_dd = gr.Dropdown(
                    label="Furry: Sub Category",
                    choices=[],
                    value=None,
                    visible=False
                )
                species_dd = gr.Dropdown(
                    label="Furry: Species",
                    choices=[],
                    value=None,
                    visible=False
                )

                gender_selection_furry = gr.Radio(
                    label="Furry Gender Selection (选择Furry性别)",
                    choices=["Male", "Female"],
                    value="Male",  # 默认公
                    visible=False  # 默认隐藏，仅在 Furry 选项时显示
                )
                
                
                def show_furry_options(opt):
                    if opt == "Trans_to_Furry":
                        return (gr.update(visible=True),
                                gr.update(visible=True),
                                gr.update(visible=True),
                                gr.update(visible=True))
                    else:
                        return (gr.update(visible=False),
                                gr.update(visible=False),
                                gr.update(visible=False),
                                gr.update(visible=False))

                gender_option.change(
                    fn=show_furry_options,
                    inputs=[gender_option],
                    outputs=[gender_selection_furry, top_cat_dd, sub_cat_dd, species_dd]
                )

                def on_top_cat_select(selected):
                    subs = get_sub_categories(FURRY_DATA, selected)
                    return gr.update(choices=subs, value=None)

                top_cat_dd.change(
                    fn=on_top_cat_select,
                    inputs=[top_cat_dd],
                    outputs=[sub_cat_dd]
                )

                def on_sub_cat_select(top_c, sub_c):
                    species = get_species_list(FURRY_DATA, top_c, sub_c)
                    return gr.update(choices=species, value=None)

                sub_cat_dd.change(
                    fn=on_sub_cat_select,
                    inputs=[top_cat_dd, sub_cat_dd],
                    outputs=[species_dd]
                )

            with gr.Column():
                user_prompt = gr.Textbox(
                    label="Original Prompt 原始提示词 (e.g. 1girl, butterfly, solo, ...)",
                    lines=5
                )
                final_output = gr.Textbox(
                    label="Transformed Output 翻译结果 (tags + description)",
                    lines=10
                )

        with gr.Row():
            translate_lang = gr.Dropdown(
                label="Translate to Language 翻译语言",
                choices=[
                    "English", "Chinese", "Japanese", "French", "German", 
                    "Italian", "Spanish", "Russian", "Dutch", "Persian", "Arabic", "Thai"
                ],
                value="English"
            )
            translated_text = gr.Textbox(
                label="Translated Result",
                lines=10
            )

        ######################################################################
        # 生成
        ######################################################################
        def on_generate(prompt, gender, gender_selection_furry, tc, sc, spc, mode, key, lang):
            # Step 1: 强制替换用户输入
            direction = transform_map.get(gender, None)
            if direction:
                prompt = forced_replace(prompt, direction)  # 替换必须在生成之前

            # Debug 替换后的 Prompt
            print(f"Debug Prompt After Replacement: {prompt}")

            # Step 2: 如果是 Furry 类目，处理性别描述
            if gender == "Trans_to_Furry":
                furry_path = f"{tc} > {sc} > {spc}" if (tc and sc and spc) else "unknown"
                species_details = FURRY_DATA.get(tc, {}).get(sc, {}).get(spc, {}).get("Details", "")
                gender_description = ", ".join([g.lower() for g in gender_selection_furry]) if gender_selection_furry else "neutral"
                prompt += f", Furry: {furry_path}, {gender_description}, {species_details}"  # 添加 Furry 类目路径

                # 添加性别描述（如果选择了性别）
                if gender_selection_furry:
                    gender_description = ", ".join([g.lower() for g in gender_selection_furry])  # 转换为小写
                    prompt += f", {gender_description}"  # 将性别描述附加到提示词
            
            # Debug 添加性别后的 Prompt
            print(f"Debug Prompt with Furry Gender: {prompt}")

            # Step 3: 提交到生成器
            merged_output = generate_transformed_output(prompt, gender, gender_selection_furry, tc, sc, spc, mode, key)

            # Debug 生成器输出
            print(f"Debug Merged Output: {merged_output}")

            # Step 4: 翻译生成的结果
            translated_output = translate_text(merged_output, lang, mode, key)

            # Step 5: 返回结果
            return merged_output, translated_output

        user_prompt.submit(
            fn=on_generate,
            inputs=[user_prompt, gender_option, gender_selection_furry, top_cat_dd, sub_cat_dd, species_dd, api_mode, api_key, translate_lang],
            outputs=[final_output, translated_text]
        )

        gen_btn = gr.Button("Generate")
        gen_btn.click(
            fn=on_generate,
            inputs=[user_prompt, gender_option, gender_selection_furry, top_cat_dd, sub_cat_dd, species_dd, api_mode, api_key, translate_lang],
            outputs=[final_output, translated_text]
        )

    return demo

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