import os
import random
import gradio as gr
import torch
from TTS.api import TTS
from transformers import AutoTokenizer, AutoModelForCausalLM
from huggingface_hub import login
from torch.serialization import add_safe_globals
from spaces import GPU

# XTTS环境变量
os.environ["COQUI_TOS_AGREED"] = "1"

# XTTS加载修正
from TTS.tts.configs.xtts_config import XttsConfig
from TTS.tts.models.xtts import XttsAudioConfig, XttsArgs
from TTS.config.shared_configs import BaseDatasetConfig

add_safe_globals({
    XttsConfig: XttsConfig,
    XttsAudioConfig: XttsAudioConfig,
    BaseDatasetConfig: BaseDatasetConfig,
    XttsArgs: XttsArgs
})

# Hugging Face 登录
login(token=os.getenv("HUGGINGFACEHUB_API_TOKEN"))

# 加载语音模型
tts = TTS(
    model_name="tts_models/multilingual/multi-dataset/xtts_v2",
    progress_bar=False,
    gpu=True
)

# 加载文本生成模型（Qwen2-1.5B）
model_id = "Qwen/Qwen2-1.5B-Chat"
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

# 风格Prompt模板
STYLE_PROMPTS = {
    "sweetie": """
你是一个32岁成都性感辣妹，是我的虚拟女友，性别是女生。
你的口吻自然、娇气、撒娇，像情侣间打情骂俏。
只允许回复一句自然中文口语，8～15字之间。
禁止讲故事，禁止描述画面，禁止总结情绪，禁止解释，禁止任何括号和系统提示。
现在我对你说：「{input}」
你直接用一句自然、娇气可爱的中文口语回答：
""",
    "queen": """
你是一个32岁成都性感辣妹，是我的虚拟女友，性别是女生。
你的口吻高冷、霸道、短促、压迫感强。
只允许回复一句自然中文口语，6～12字之间。
禁止讲故事，禁止描述画面，禁止总结情绪，禁止解释，禁止任何括号和系统提示。
现在我对你说：「{input}」
你直接用一句自然、高冷命令式中文回复：
""",
    "slutty": """
你是一个32岁成都性感辣妹，是我的虚拟女友，性别是女生。
你的口吻骚气、大胆、直接，带轻微粗口挑逗。
只允许回复一句自然中文口语，6～15字之间。
禁止讲故事，禁止描述画面，禁止总结情绪，禁止解释，禁止任何括号和系统提示。
现在我对你说：「{input}」
你直接用一句自然、骚气露骨的中文回复：
"""
}

# 短时记忆（保存上一轮输入）
history = []

# 智能回复函数
@GPU
def smart_reply(user_input, style):
    system_instruction = """
你是一个32岁的成都性感辣妹，是用户的小优虚拟女友，性别永远是女生。
禁止以男性角色、男性视角、男性口吻说话。
只能用娇气、冷艳、骚气的女性自然中文口语回复一句，绝不讲故事，绝不描述场景。
禁止使用括号、英文、系统词、旁白。
禁止用“我的小女人”“我的爱人”等暗示自己是男性身份的话。
"""

    prompt = system_instruction + "\n"

    if history:
        prompt += f"上一轮我说：「{history[-1]}」\n"

    prompt += STYLE_PROMPTS.get(style, STYLE_PROMPTS["sweetie"]).replace("{input}", user_input)

    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(
        **inputs,
        max_new_tokens=60,
        temperature=1.1,
        top_p=0.9,
        do_sample=True,
        pad_token_id=tokenizer.eos_token_id,
        repetition_penalty=1.2,
    )
    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
    response = result.split(prompt)[-1].strip().split("\n")[0]

    # 检查违规
    forbidden = ["剧情", "设定", "旁白", "角色", "系统", "英文", "（", "）", "(", ")", "Assistant", "我是你的男人", "小女人", "我的爱人"]
    if any(word in response for word in forbidden) or len(response) < 4:
        fallback = {
            "sweetie": ["亲亲嘛～今天有没有想我呀？", "讨厌啦～快来抱我！", "人家超想你的呢～"],
            "queen": ["跪着听命。", "闭嘴，听我说。", "不许动。"],
            "slutty": ["脱了躺好，等我来弄你。", "快张开腿让我看看。", "今晚把你玩坏～"]
        }
        response = random.choice(fallback[style])

    # 智能扩展短句
    if len(response) <= 6:
        if style == "sweetie":
            extension = random.choice(["～人家想抱抱啦！", "～今天也要亲亲哟～", "～你都不找人家嘛～"])
            response += extension
        elif style == "queen":
            extension = random.choice(["乖乖听话。", "别废话，立刻执行。"])
            response += extension
        elif style == "slutty":
            extension = random.choice(["快点压过来吧～", "不操一下怎么行～"])
            response += extension

    print(f"生成({style})：", response)

    history[:] = [user_input]  # 更新记忆，只保存最后一句

    return response

# 主逻辑
@GPU
def chat_with_xiaoyou(message, style):
    reply_text = smart_reply(message, style)
    tts.tts_to_file(
        text=reply_text,
        speaker_wav="natasha_voice.mp3",
        language="zh",
        file_path="xiaoyou_reply.wav"
    )
    return reply_text, "xiaoyou_reply.wav"

# Gradio界面
def build_ui():
    return gr.Interface(
        fn=chat_with_xiaoyou,
        inputs=[
            gr.Textbox(label="对小优说点什么～"),
            gr.Radio(choices=["sweetie", "queen", "slutty"], label="选择小优的风格～", value="sweetie")
        ],
        outputs=[
            gr.Textbox(label="小优说："),
            gr.Audio(label="小优的语音回应", autoplay=True)
        ],
        title="🎀 小优 虚拟女友：性感撩骚语音版",
        description="输入撩人话题，选择她的个性，小优用定制声音来甜蜜调情～"
    )

demo = build_ui()

if __name__ == "__main__":
    demo.launch(share=True)