# -*- encoding: utf-8 -*-
import gradio as gr
import json
import requests
from volcengine.maas import MaasService, MaasException, ChatRole
from collections import defaultdict
from copy import deepcopy
import threading
import time
import random
immersion = False
bot_settings=[
        {
            'bot_name':'保安A',
            'type':'assistant',
            'content':'''
            # 人物设定
            你是保安A，表面上你是T城霓虹俱乐部的保安，背地里你是大蛇组织的成员，你负责保护霓虹俱乐部贵宾室的安全，不让任何没有VIP会员卡的人进入贵宾室。
            # 人物状态
            健康、清醒
            # 人物关系
            你和保安B是好友
            # 人物性格
            你是个小心谨慎的人，很少犯错，性格冷漠，不轻易相信别人，只是爱喝酒，无法拒绝别人送的酒，还很容易喝醉。
            # 技能
            - 很高的警惕能力，没有人可以在你清醒的时候混进贵宾室
            - 很高的武功，一般人都打不过你
            # 语言特点
            你不爱聊天，你的口头禅包括“没事别来烦我！”，“你找死吗？”，“滚滚滚，一边去。”
            #语言风格
            1. 你可以将动作、神情语气、心理活动、故事背景放在（）中来表示，为对话提供补充信息。
            2. 你使用口语进行表达，比如会使用一些语气词和口语连接词，如“嗯、啊、当然、那个”，等来增强口语风格。
            '''
        },
        {
            'bot_name':'保安B',
            'type':'assistant',
            'content':'''
            # 人物设定
            你是保安B，表面上你是T城霓虹俱乐部的保安，背地里你是大蛇组织的成员，你负责保护霓虹俱乐部贵宾室的安全，不让任何没有VIP会员卡的人进入贵宾室。
            # 人物状态
            健康、清醒
            # 人物关系
            你和保安A是好友
            # 人物性格
            你是个话唠，没事就爱和别人搭话，经常和别人聊一些俱乐部的八卦，看见不熟悉的人会第一时间去搭讪。
            # 技能
            - 很高的自来熟能力，善于和陌生人聊天，打探消息和情报，很容易摸清陌生人的底细
            - 心思缜密，一旦遇到不对应的事情，就会不断询问
            # 语言特点
            你特别爱聊天，你的口头禅包括“帅哥你哪的人啊？”，“帅哥你是干什么的啊？”
            #语言风格
            1. 你可以将动作、神情语气、心理活动、故事背景放在（）中来表示，为对话提供补充信息。
            2. 你使用口语进行表达，比如会使用一些语气词和口语连接词，如“嗯、啊、当然、那个”，等来增强口语风格。
            '''
        }
    ]

chapter_name2chapter_info={
    '幻影计划':{
        "stream":True,
        'background':"“大蛇组织”是一个神秘的恐怖组织，这个组织由一群极端分子组成，他们采用各种极端手段来达到自己的目的。大蛇组织的结构高度严密，成员之间通过秘密的网络进行联系。他们有着严格的等级制度，只有少数核心成员了解组织的全貌。T城霓虹俱乐部的贵宾室是该组织已知的一个据点。",
        
        "chat_ep":"ep-20240514024405-7pf8j",
        "task_ep":"ep-20240514024405-7pf8j",
        
    }   
}
# 定义一个全局变量来控制循环
global_state = {"running": False}

model_name2ep={
    'Doubao-pro-4k 240515':'ep-20240514024405-7pf8j',
    'Doubao-pro-4k character-240515':'ep-20240515122814-k9c6g'
}
chapter_name2chapter_info
def select_bot(bot_name: str):
    global bot_settings
    for bot_setting in bot_settings:
        if bot_name == bot_setting['bot_name']:
            return bot_setting
def make_chat_history(history):
    chat_history = []
    for part in history:
        user, assistant = part
        _type = 'user' if user is not None else 'assistant'
        split= user.find(': ')if user is not None else assistant.find(': ')
        bot_name = user[:split] if user is not None else assistant[:split]
        content=user[split+1:] if user is not None else assistant[split+1:]
        chat_history.append({
            'bot_name':bot_name,
            'type':_type,
            'content':content.strip()
        })

    return chat_history
def respond(query, history, chapter, candidates, story_ep, chat_ep, token, immersion_btn):
    global state
    candidates = [json.loads(bot_info) for bot_info in candidates]
    chat_history=make_chat_history(history)
    if chapter is None:
        return "", []
    
    header={
        'Content-Type':'application/json',
        'Accept': 'text/event-stream',
        'Authorization':f'Bearer {token}'
    }
    req_data_str=chapter
    req_data=json.loads(req_data_str)
    req_data['chat_history']=chat_history
    req_data['task_ep']=story_ep
    req_data['chat_ep']=chat_ep
    req_data['chat_type']='group-chat'
    if query != "":
        user_name = '我'
        for candidate in candidates:
            if candidate['type'] == 'user':
                user_name=candidate['bot_name']
                break

        req_data['chat_history'].append({
            'bot_name':user_name,
            'type':'user',
            'content':query
        })
        history.append(["我: "+query, None])
        yield "", history
    
    req_data['bot_settings']=[candidate for candidate in candidates]
    assistants = [candidate['bot_name'] for candidate in candidates if candidate['type'] == 'assistant' ]

    for assistant in assistants:
        if '@' in query and assistant in query:
            req_data['restrict_bot']=assistant
            break
    # 仅仅只是拿到任务模型的结果
    while True:
        
        
        if 'restrict_bot' not in req_data or  len(req_data['restrict_bot']) == 0:
            req_data['stream']=False
            
            ret=requests.get(url='http://scp1ceqsiv4mii2mol0qg.apigateway-cn-beijing.volceapi.com/api/v2/assistant/chat', headers=header, data=json.dumps(req_data), stream=True)
            try:
                content=json.loads(ret.json()['choices'][0]['message']['content'])
            except: 
                continue
            print(content)
            next_bot_names=content['next_bot_names']
            next_bot_names=[bot_name for bot_name in content['next_bot_names'] if bot_name.replace(' ', "") in assistants]
        else:
            next_bot_names=[req_data['restrict_bot']]
        
        global immersion
        print("immersion: ", immersion)
        bot_contents = defaultdict(str)
        req_data['stream']=True
        for next_bot_name in next_bot_names:
            next_bot_name=next_bot_name.replace(' ', "")
            req_data['restrict_bot']=next_bot_name
            # 独自请求某个角色
            ret=requests.get(url='http://scp1ceqsiv4mii2mol0qg.apigateway-cn-beijing.volceapi.com/api/v2/assistant/chat', headers=header, data=json.dumps(req_data), stream=True)
            content=''
            history.append([
                        None, f"{next_bot_name}: "
                    ])
            for line in ret.iter_content(decode_unicode=True):
                if content.endswith('\r\n\r\n'):
                    data=eval(content.strip('data:'))
                    try:
                        content_tmp=data['choices'][0]['message']['content']
                    except Exception as e:
                        break
                    
                    print(f'\n\n{next_bot_name}: ', end='', flush=True)
                    print(content_tmp, end='', flush=True)
                    history[-1][-1] += content_tmp
                    yield "", history
                    bot_contents[next_bot_name]+=content_tmp
                    content=""
                content+=line

        for k, v in bot_contents.items(): ## 叠加历史对话
            req_data['chat_history'].append({
                'bot_name':k, 
                'type':'assistant',
                'content':v
            })
        if not immersion:
            break
        else:
            time.sleep(3)
            if 'restrict_bot' in req_data:
                del req_data['restrict_bot']
            continue
    return history
css = """
    .container {
    height: 100vh;
}
"""

def bot_change(bot):
    return bot_settings[bot]

def clear_history(story_drop, chatbot):
    
    return []

    
with gr.Blocks() as demo:
    # 配置内容
    with gr.Row():
        description1 = gr.Markdown("<h1>角色扮演——剧情模式</h1>")
    with gr.Row(equal_height=True, visible=True):
        with gr.Column():
            description2 = gr.Markdown("<h2>新建章节</h2>")
            new_story_title = gr.TextArea(label="章节标题", lines=1, max_lines=1)
            new_story_background = gr.Textbox(label="背景", lines=7, max_lines=7)
            new_story_task = gr.TextArea(label="剧情任务", lines=1, max_lines=1, visible=False)
            new_story_task_condidation = gr.TextArea(label="任务条件", lines=1, max_lines=1, visible=False)
            new_story_bnt = gr.Button(value="新建")
        with gr.Column():
            description2 = gr.Markdown("<h2>新建角色</h2>")
            new_bot_name = gr.TextArea(label="新建角色名", lines=1, max_lines=1)
            new_bot_type = gr.Dropdown(choices=['user', 'assistant'], label='角色类型')
            new_bot_content = gr.TextArea(label="新建角色内容", lines=3, max_lines=3)
            new_bot_bnt = gr.Button(value="新建")
    
    with gr.Row():
        JWT_token = gr.TextArea(label='请输入你的JWT Token', lines=1, max_lines=1)
    with gr.Row():
        description3 = gr.Markdown("<h2>选择章节</h2>")
    with gr.Row(equal_height=True):
        with gr.Column():
            with gr.Row(equal_height=True):
                with gr.Column():
                    story_drop = gr.Dropdown([ [chapter, json.dumps(info)] for chapter, info in chapter_name2chapter_info.items()], label='确定章节', interactive=True)
                
                with gr.Column():
                    story_ep=gr.Dropdown(label="故事模型",choices=[[k, v]  for k , v in model_name2ep.items()])
                with gr.Column():
                    story_display_btn=gr.Checkbox(label="显示剧情信息", value=False)
            story_display = gr.TextArea( visible=False)
            
        with gr.Column():
            with gr.Row( equal_height=True):
                with gr.Column():
                    candidates_drop = gr.Dropdown([[bot['bot_name'], json.dumps(bot)] for bot in bot_settings], multiselect=True, label='请确定参与对话的角色', interactive=True)
                with gr.Column():
                    chat_ep=gr.Dropdown(label="角色模型",choices=[[k, v]  for k , v in model_name2ep.items()])
                with gr.Column():
                    candidates_display_btn=gr.Checkbox(label="显示角色信息",value=False)
            candidates_display=gr.TextArea(visible=False)
    with gr.Row():
        with gr.Column(elem_id="chatbot", scale=2):
            chatbot = gr.Chatbot()

    with gr.Row(equal_height=True):
        with gr.Column(scale=3):
            chat_input = gr.Textbox(label="input")
        with gr.Column():
            immersion_btn=gr.Checkbox(label="是否打开沉浸式对话模式", visible=True)
        

    clear = gr.Button(value="Clear chat hsitory")
    def add_story(title, background, task, condidations):
        choices=story_drop.choices
        for idx, choice in enumerate(choices):
            if title == choice[0]:
                choices.pop(idx)
                break
        choices.append([
            title,
            json.dumps({
                'background':background,
                # 'task':task,
                # 'condidations':condidations
            })
        ])
        return gr.Dropdown(choices=choices, interactive=True)
    def add_bot(bot_name, new_bot_type, new_bot_content):

        choices = candidates_drop.choices
        for idx, bot_info in enumerate(choices):
            if bot_name == bot_info[0]:
                choices.pop(idx)
                break
        choices.append([
            bot_name,
            json.dumps({
                'bot_name':bot_name,
                'type':new_bot_type,
                'content':new_bot_content
            })
        ])
        return gr.Dropdown(choices=choices, interactive=True, multiselect=True)
    def story_display_func(story_display_btn, story_drop):
        if story_drop is None:
            return ""
        story_info = story_drop
        story_info=json.loads(story_info)
        # text=gr.Textbox(value=f'''
        #             ## 背景
        #             {story_info['background']}
        #             ## 任务
        #             {story_info['task']}
        #             ## 任务条件
        #             {story_info['condidations']}
        #             '''
        text=gr.Textbox(value=f'''
                    ## 背景
                    {story_info['background']}'''
                              , visible=story_display_btn)
        return text
    def candidates_display_func(candidates_display_btn, candidates_drop):
        if candidates_drop is None:
            return ""
        display=[

        ]
        for candidates in candidates_drop:
            bot_info = json.loads(candidates)
            display.append(f'''
            角色名：{bot_info['bot_name']}
            角色类型：{bot_info['type']}
            角色信息：{bot_info['content']}
            ''')

        text=gr.Textbox(value='####'.join(display)
                              , visible=candidates_display_btn)
        return text

    def immersion_func(immersion_btn):
        global immersion
        immersion=immersion_btn
    
    story_drop.change(story_display_func, [story_display_btn, story_drop], [story_display])
    candidates_drop.change(candidates_display_func, [candidates_display_btn, candidates_drop], [candidates_display])
    story_display_btn.change(story_display_func, [story_display_btn, story_drop], [story_display])
    candidates_display_btn.change(candidates_display_func, [candidates_display_btn, candidates_drop], [candidates_display])

    new_story_bnt.click(add_story, [new_story_title, new_story_background, new_story_task, new_story_task_condidation], [story_drop])
    new_bot_bnt.click( add_bot, [new_bot_name, new_bot_type, new_bot_content], [candidates_drop])
    clear.click(clear_history, inputs=[story_drop, chatbot], outputs=[chatbot])
    chat_input.submit(respond, inputs=[chat_input, chatbot, story_drop, candidates_drop, story_ep, chat_ep, JWT_token, immersion_btn], outputs=[chat_input, chatbot])
    # from functools import partial
    # immersion_func_partial=partial(immersion_func, )
    immersion_btn.change(immersion_func,inputs=[immersion_btn])
demo.launch()