import base64
import datetime
from hashlib import md5
import time
import requests
import uuid
import json
import traceback
import sys

# 测试环境
base_url = "https://ai-agent-test.7moor.com/agent"
# 生产环境
# base_url = "https://ai-agent-sdk.7moor.com/agent"

# 填写对应的APPID，应用ID
APPID = ""
# 填写对应的APPTOKEN，应用令牌
APPTOKEN = ""
# 填写对应的AGENT_ID，agent编号
AGENT_ID = ""


def get_auth_token(appid, timestamp):
    """
    生成认证令牌

    Args:
        appid (str): 应用ID
        timestamp (str): 时间戳

    Returns:
        str: Base64编码的认证令牌
    """
    auth_token = base64.b64encode(f'{appid}:{timestamp}'.encode('utf-8')).decode('utf-8')
    return auth_token


def get_auth_sig(appid, apptoken, timestamp):
    """
    生成认证签名

    Args:
        appid (str): 应用ID
        apptoken (str): 应用令牌
        timestamp (str): 时间戳

    Returns:
        str: MD5加密的签名字符串
    """
    sig = md5(f'{appid}{apptoken}{timestamp}'.encode('utf-8')).hexdigest()
    return sig


def init_agent(sid, timestamp):
    """
    初始化会话

    Args:
        sid (str): 会话ID
        timestamp (str): 时间戳

    Returns:
        dict: 会话初始化响应结果
    """
    try:
        auth_token = get_auth_token(APPID, timestamp)
        sig = get_auth_sig(APPID, APPTOKEN, timestamp)

        payload = {
            "appid": APPID,
            "apptoken": APPTOKEN,
            "timestamp": timestamp,
            "auth_token": auth_token,
            "sig": sig,
            "sid": sid,
            "agent_id": AGENT_ID,
            "msg_receive_url": "http://chat-task-based-addon/ac/callback",
            "session_source": "autotest",
            "system_args": {}
        }
        req_url = base_url + "/v1/api/init_session"
        headers = {
            "Content-Type": "application/json",
        }
        response = requests.request("POST", req_url, headers=headers, json=payload)
        return json.loads(response.text)
    except Exception as e:
        print(f"[{sid}] exception:", traceback.format_exc())
        return {"error": "网络异常！请重试"}


def query_with_history(sid, messages, timestamp):
    """
    与AI AGENT进行对话

    Args:
        sid (str): 会话ID
        messages (list): 对话消息列表
        timestamp (str): 时间戳

    Returns:
        str: AI的完整回复内容
    """
    try:
        start_time = datetime.datetime.now()
        auth_token = get_auth_token(APPID, timestamp)
        sig = get_auth_sig(APPID, APPTOKEN, timestamp)

        req_url = base_url + "/v1/api/query_with_history"
        payload = {
            "appid": APPID,
            "apptoken": APPTOKEN,
            "timestamp": timestamp,
            "auth_token": auth_token,
            "sig": sig,
            "sid": sid,
            "agent_id": AGENT_ID,
            "messages": messages,
            "stream": True
        }
        headers = {
            "Content-Type": "application/json",
        }
        response = requests.request("POST", req_url, headers=headers, json=payload, timeout=120)

        full_reply = ""  # 用于存储完整回复
        for chunk in response.iter_lines():
            if chunk:
                chunk_str = chunk.decode("utf-8")
                print(chunk_str)
                chunk_obj = json.loads(chunk_str)

                # 获取当前块的回复内容
                content = chunk_obj.get("reply", "")
                if content:
                    full_reply += content
                    # 实时打印当前内容
                    for char in content:
                        sys.stdout.write(char)
                        sys.stdout.flush()
                        # time.sleep(0.01)

                if chunk_obj.get("is_end") == True:
                    end_time = datetime.datetime.now()
                    print()  # 换行
                    print(f"[{sid}] 总耗时:{end_time - start_time}")
                    return full_reply

    except Exception as e:
        print(f"[{sid}] exception:", traceback.format_exc())
        return "网络异常！请重试"


def close_agent(sid, timestamp):
    """
    关闭会话

    Args:
        sid (str): 会话ID
        timestamp (str): 时间戳

    Returns:
        dict: 关闭会话响应结果
    """
    try:
        auth_token = get_auth_token(APPID, timestamp)
        sig = get_auth_sig(APPID, APPTOKEN, timestamp)

        payload = {
            "appid": APPID,
            "apptoken": APPTOKEN,
            "timestamp": timestamp,
            "auth_token": auth_token,
            "sig": sig,
            "sid": f"{sid}"
        }
        req_url = base_url + "/v1/api/close_session"
        headers = {
            "Content-Type": "application/json"
        }
        response = requests.request("POST", req_url, headers=headers, json=payload)
        print(json.loads(response.text))
    except Exception as e:
        print(f"[{sid}] exception:", traceback.format_exc())
        return {"error": "网络异常！请重试"}


if __name__ == '__main__':
    # 生成会话ID
    sid = uuid.uuid4().hex
    # 生成时间戳
    timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
    # timestamp = '20240317093513'
    print(f"当前会话sid:{sid}，timestamp:{timestamp}")

    # 1、初始化会话
    init_agent(sid, timestamp)

    # 存储对话历史
    messages = []
    try:
        while True:
            user_input = input("你: ")
            if user_input.lower() in ["exit", "quit", "bye"]:
                break

            messages.append({
                "type": "user",
                "text": user_input
            })

            print("agent: ", end="")  # 提示AI回复开始
            response = query_with_history(sid, messages, timestamp)

            messages.append({
                "type": "assistant",
                "text": response
            })
    except KeyboardInterrupt:
        print("用户中断会话")
    finally:
        # 确保关闭会话
        close_agent(sid, timestamp)