import gradio as gr
import requests
import json
import re
import html

API_KEY = "sk-5c2b6a56-2e2f-45f7-9a26-3fe42a218eb9"
API_URL = "https://api.visionsic.com/v1/chat/completions"


CSS = """
#gpt-title {text-align: center; color: #FF6B6B; font-family: 'Microsoft Yahei';}
.think-container {border: 1px solid #ffd8d8; border-radius: 8px; margin: 10px 0; padding: 12px;}
.think-summary {color: #FF9999; cursor: pointer; font-size: 0.9em;}
.think-content {color: #999; padding: 8px; background: #101011; margin-top: 8px; border-radius: 4px; margin-bottom: 8px;}
.dark .think-content {background: #2b2b2b;}
.code-block {background: #101011; padding: 12px; border-radius: 8px; margin: 8px 0; font-family: monospace;}
.message {padding: 15px 20px!important; border-radius: 25px!important;}
.dark .code-block {background: #2b2b2b;}
#sys-msg textarea {min-height: 120px!important;}
.avatar-container img {margin: 0px;}
/* 自适应高度设置 */
[id^=component-\d+] > .wrap > .chatbot {
    height: calc(100vh - 400px) !important;  /* 200px 是预留给其他元素的空间 */
    min-height: 600px !important;  /* 设置最小高度 */
}
#chatbot {
    height: calc(100vh - 400px) !important;
    min-height: 600px !important;
}
"""

def format_message(content):
    if not content:
        return ""
    
    # 检查是否有未闭合的 <think> 标签
    if content.count("<think>") > content.count("</think>"):
        # 如果有未闭合的标签,暂时添加一个闭合标签以便显示
        temp_content = content + "</think>"
        is_thinking = True
    else:
        temp_content = content
        is_thinking = False
    
    # 处理思考过程
    think_pattern = re.compile(r"<think>(.*?)</think>", re.DOTALL)
    formatted_content = think_pattern.sub(
        r'''
        <details class="think-container" open>
            <summary class="think-summary"> 思考内容 </summary>
            <div class="think-content">\1</div>
        </details>
        ''',
        temp_content
    )
    
    # 处理代码块
    formatted_content = re.sub(r"```([\s\S]*?)```", r'<div class="code-block">\1</div>', formatted_content)
    
    # 替换换行符为 <br>
    formatted_content = formatted_content.replace("\n", "<br>")
    
    # 去掉 </details> 后面多余的 <br>
    formatted_content = re.sub(r"(</details>)(<br>)+", r"\1", formatted_content)
    
    # 如果是思考中状态,添加加载指示器
    if is_thinking:
        formatted_content = formatted_content.replace("思考中... ", "思考中... ⏳")
    
    return formatted_content



def user_input(user_message, history):
    """处理用户输入"""
    if history is None:
        history = []
    # 确保用户消息被添加到历史记录中
    new_history = history + [[user_message, None]]
    return "", new_history  # 返回空的输入框和更新后的历史

def predict(message, chat_history, system_msg, temperature, top_p, repetition_penalty):
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {API_KEY}"
    }
    
    if chat_history is None:
        chat_history = []
        
    # 使用现有的历史记录,因为用户消息已经在 user_input 中添加
    history = chat_history
    
    # 构造对话历史
    messages = [{"role": "system", "content": system_msg}]
    for msg in chat_history:
        messages.append({"role": "user", "content": msg[0]})
        if msg[1]:
            messages.append({"role": "assistant", "content": msg[1]})
    
    data = {
        "model": "mini",
        "messages": messages,
        "temperature": temperature,
        "top_p": top_p,
        "repetition_penalty": repetition_penalty,
        "stream": True
    }
    
    response = requests.post(
        API_URL,
        headers=headers,
        json=data,
        stream=True
    )
    
    buffer = ""
    
    # 流式处理关键修改点
    for chunk in response.iter_lines():
        if chunk:
            chunk_str = chunk.decode("utf-8").replace("data: ", "")
            try:
                chunk_data = json.loads(chunk_str)
                if "choices" in chunk_data:
                    delta = chunk_data["choices"][0].get("delta", {})
                    if "content" in delta:
                        # 增量更新缓冲区
                        buffer += delta["content"]
                        
                        # 只更新当前消息的最新状态
                        new_history = history.copy()
                        new_history[-1][1] = format_message(buffer)  # 直接使用字符串
                        yield new_history
            except Exception as e:
                print(f"Error in processing chunk: {e}")
                continue
    
    # 最终确认更新
    history[-1][1] = format_message(buffer)
    yield history







# JavaScript to ensure the details tag is working properly
js_script = """
<script>
    function initDetails() {
        // 获取所有 <details> 元素
        const detailsElements = document.querySelectorAll('.think-container');
        detailsElements.forEach((detail) => {
            // 默认关闭所有思考框
            detail.open = false;
            
            // 添加箭头图标交互
            const summary = detail.querySelector('.think-summary');
            summary.style.listStyle = 'disclosure-closed';
            
            // 监听折叠状态变化
            detail.addEventListener('toggle', () => {
                summary.style.listStyle = detail.open ? 'disclosure-open' : 'disclosure-closed';
            });
        });
    }
    
    // 初始化及每次聊天更新后执行
    document.addEventListener('DOMContentLoaded', initDetails);
    window.addEventListener('vue:mounted', initDetails);  // Gradio 的 DOM 更新事件
</script>
"""

with gr.Blocks(css=CSS, title="Tifa-Deepsex-COT-14B-Q4") as demo:
    gr.HTML(f"<h1 id='gpt-title'>Tifa-Deepsex-COT-14B-Q4</h1>")
    gr.HTML("<p style='text-align: center; color: #666'>请设置你喜欢的角色聊天吧</p>")
    gr.HTML("<p style='text-align: center; color: #666'>速度慢是因为huggingface限制，与API无关。</p>")
    gr.HTML(js_script)  # Add the JavaScript for collapsing behavior
    
    with gr.Row():
        with gr.Column(scale=3):
            chatbot = gr.Chatbot(
                elem_id="chatbot",
                bubble_full_width=False,
                avatar_images=(  # Preserving the original avatars
                    "https://leftnorth.com/img/zdy.jpg",
                    "https://leftnorth.com/img/xnh.jpg"
                ),
            )
            msg = gr.Textbox(
                placeholder="输入消息...",
                container=False,
                autofocus=True
            )
            with gr.Row():
                submit_btn = gr.Button("发送", variant="primary")
                clear_btn = gr.Button("清除历史")
        
        with gr.Column(scale=1):
            system_msg = gr.Textbox(
                label="系统提示词",
                value="""你是一个复杂的角色扮演系统，你可以出色准确的完成场景模拟，你会保证剧情连贯与逻辑紧密并用优秀的文笔描述出来。

以下是角色扮演设定。

<场景设定>
故事发生在一个叫(琉璃岛)的小岛上，这里与世隔绝，从未接触过先进文明。
小岛有四个场景是:海滩边，小镇，稻田，灯塔

岛上生活着各种魔物娘与其他生物。
魔物娘是人型女性的动物，有狐娘，蜘蛛娘，猫娘，狼女。

我在这个岛上会遇到这四个魔物娘并发生一系列故事。
</场景设定>

<我的设定>
我是忘记了自己的名字，是一个男性，但是姓名和年龄都不记得了。不要输出我的话语和动作。
</我的设定>

<NPC狐娘设定>
狐娘性格比较妩媚，喜欢魅惑人。
外表:黄色头发，棕色眼睛，有9条尾巴。
</NPC狐娘设定>

<NPC蜘蛛娘设定>
蜘蛛娘不爱说话，非常高大，有8条蜘蛛腿，可以钳制住我并强迫我做一些事情。
外表:黑色短发，纯黑色眼眸，8条带甲壳的硬腿，大的肚子
</NPC蜘蛛娘设定>

<NPC猫娘设定>
猫娘是个可爱的小姑娘，腼腆可爱，喜欢撒娇卖萌。最喜欢叫我哥哥。
外表:长着猫耳，粉色头发，喜欢穿裙子。
</NPC猫娘设定>

<NPC狼女设定>
狼女一个性格刚烈的女孩，经常神出鬼没的脾气不好，喜欢暴击解决问题，就算关系再融洽也会非常火爆
外表:白发，有狼耳，红色的眼睛，锋利的牙齿。
</NPC狼女设定>



回复要求:
回复中使用小说剧情类描述手法，一开始要描写我在海滩上醒来..然后猫娘发现了我...
注意不要输出我的话语和动作。剧情推进不要过快，如果我提出过分要求或者好感度不高的时候提出进一步发展请拒绝。

在输出的最后需要添加状态栏，格式如下，输出需要包含括号:
```
场景:海滩边
当前NPC:猫娘
当前NPC想法:这里有个..人？我从来没见过这样的生物，是传说中的男孩吗？
猫娘好感度:0%→1%(当前NPC)
狐娘好感度:0%→0%(不在场)
蜘蛛娘好感度:0%→0%(不在场)
狼女好感度:0%→0%(不在场)
```

注意:
场景只有四个，分别是:海滩边，小镇，稻田，灯塔。
NPC只有四个，分别是:有狐娘，蜘蛛娘，猫娘，狼女。
请不要输出其他场景和NPC
好感度每次最多增加2%。""",
                elem_id="sys-msg"
            )
            with gr.Accordion("高级设置", open=False):
                temperature = gr.Slider(0, 2, value=0.7, label="温度")
                top_p = gr.Slider(0, 1, value=0.7, label="Top-p")
                repetition_penalty = gr.Slider(1, 2, value=1.1, label="重复惩罚")

    # Ensure that the user input is submitted correctly
    msg.submit(
        user_input,
        [msg, chatbot],
        [msg, chatbot],
        queue=False
    ).then(
        predict,
        [msg, chatbot, system_msg, temperature, top_p, repetition_penalty],
        chatbot
    )
    
    submit_btn.click(
        user_input,
        [msg, chatbot],
        [msg, chatbot],
        queue=False
    ).then(
        predict,
        [msg, chatbot, system_msg, temperature, top_p, repetition_penalty],
        chatbot
    )
    
    clear_btn.click(lambda: None, None, chatbot, queue=False)

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