from flask import request
import os, sys

base_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(base_path)
from ai_chat_stream import parse_utils, sys_content,parse_comment_utils
from ai_utils import get_bot_with_history_stream, get_ai_chat_history, get_bot_with_history_rag_stream, \
    get_ai_chat_history_stream
from parse_utils import get_last_user_value, convert_dialog_format, create_res_data,extract_user_contents
from flask import Flask, Response
import json

app = Flask(__name__)


def sse_generator(data):
    res_data = create_res_data("正在识别用户问题", 3)
    yield f"data: {json.dumps(res_data, ensure_ascii=False)}\n\n".encode('utf-8')
    sys_list = [{"role": "system", "content": sys_content.action_content}]
    ai_reply, total_tokens = get_ai_chat_history(sys_list + get_last_user_value(data))

    if ai_reply.get('content') == '1':
        res_data = create_res_data("正在解析链接", 3)
        yield f"data: {json.dumps(res_data,ensure_ascii=False)}\n\n".encode('utf-8')
        convert_dialog_format_text = '拼接'.join(extract_user_contents(convert_dialog_format(data)))
        yield from parse_utils.extract_mp4_to_text(convert_dialog_format_text)

    elif ai_reply.get('content') == '2':
        res_data = create_res_data("正在文案改写", 3)
        yield f"data: {json.dumps(res_data, ensure_ascii=False)}\n\n".encode('utf-8')
        sys_list = [{"role": "system", "content": sys_content.douyin_text_fix}]
        bot_result_list = get_ai_chat_history_stream(sys_list + convert_dialog_format(data))
        if bot_result_list:
            for result in bot_result_list:
                yield result
        res_data = create_res_data("WRITE", 10)
        yield f"data: {json.dumps(res_data, ensure_ascii=False)}\n\n".encode('utf-8')

    elif ai_reply.get('content') == '3':
        res_data = create_res_data("正在查询爱搜知识库", 3)
        yield f"data: {json.dumps(res_data, ensure_ascii=False)}\n\n".encode('utf-8')
        sys_list = [{"role": "system", "content": sys_content.helpful_content_text_rag}]
        bot_result_list = get_bot_with_history_rag_stream(sys_list + convert_dialog_format(data))
        if bot_result_list:
            for result in bot_result_list:
                yield result


    elif ai_reply.get('content') == '4':
        res_data = create_res_data("正在获取评论", 3)
        yield f"data: {json.dumps(res_data, ensure_ascii=False)}\n\n".encode('utf-8')
        convert_dialog_format_text = '拼接'.join(extract_user_contents(convert_dialog_format(data)))
        yield from parse_comment_utils.extract_comment_to_text(convert_dialog_format_text)
        res_data = create_res_data("WRITE,VIDEO", 10)
        yield f"data: {json.dumps(res_data, ensure_ascii=False)}\n\n".encode('utf-8')

    else:
        res_data = create_res_data("正在联网搜索", 3)
        yield f"data: {json.dumps(res_data, ensure_ascii=False)}\n\n".encode('utf-8')
        sys_list = [{"role": "system", "content": sys_content.helpful_content_text}]
        bot_result_list = get_bot_with_history_stream(sys_list + convert_dialog_format(data))
        if bot_result_list:
            for result in bot_result_list:
                yield result

@app.route('/api/chat', methods=['POST'])
def sse_stream():
    data = request.get_json()
    return Response(
        sse_generator(data),
        mimetype="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no"
        }
    )


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=8087)
