from openai import OpenAI
import sys
import os
import threading
import time

# 温馨提示:
# 1.需要登录https://platform.deepseek.com/api_keys 创建api并充值
# 2.增加环境变量
# linux&mac 示例:
#   修改~/.bashrc 或~/.zshrc 添加:
#   export DEEPSEEK_API_KEY="sk-xxxxxxxxxxxxxxx"
#   source ~/.zshrc
api_key = os.getenv('DEEPSEEK_API_KEY')

global_msg = []

client = OpenAI(api_key=api_key, base_url="https://api.deepseek.com")

# 控制 loading 动画的全局变量
loading = False

def loading_animation():
    """
    显示 loading 动画
    """
    while loading:
        for symbol in "|/-\\":
            print(f"\r>> {symbol} 正在思考...", end="", flush=True)
            time.sleep(0.1)
    print("\r>> ", end="", flush=True)  # 清除 loading 动画

def get_reply(input_msg):
    global loading
    msg = {"role": "user", "content": input_msg}
    global_msg.append(msg)

    # 启动 loading 动画
    loading = True
    loading_thread = threading.Thread(target=loading_animation)
    loading_thread.start()

    response = client.chat.completions.create(
        model="deepseek-chat",
        messages=global_msg,
        stream=True
    )
    # 流式输出
    reply = ""
    for chunk in response:
        if chunk.choices[0].delta.content:
            if loading:
                loading = False  # 停止 loading 动画
                loading_thread.join()  # 等待 loading 线程结束
            reply += chunk.choices[0].delta.content
            print(chunk.choices[0].delta.content, end='', flush=True)

    # 存储多轮对话
    global_msg.append({"role": "assistant", "content": reply})
    return reply


def print_banner():
    banner = r'''
      _                        _     
     | |                      | |    
     | |__  _ __ ___   ___ ___| |__  
     | '_ \| '_ ` _ \ / __/ __| '_ \ 
     | | | | | | | | | (_| (__| |_) |
     |_| |_|_| |_| |_|\___\___|_.__/ 
                                     
     这是一个小demo，用来和AI对话. 支持会话记忆，你可以设定AI的角色，然后提问。
     举例:
         >>请将我下面说的每一句话翻译成英文，尽量翻译得优美一些。
         >>当然，我很乐意帮您翻译。请告诉我您想翻译的句子，我会尽量用优美、地道的英语表达出来。
         >>今天天气不错，适合户外骑行。
         >>The weather today is splendid, perfect for a cycling adventure in the embrace of nature.
         >>白毛浮绿水，红掌拨清波。
         >>White feathers float upon the emerald waters, while red palms gently part the crystal waves.
         >>举头望明月，低头思故乡。
         >>I lift my gaze to the moon, a celestial lantern in the night, and bow my head, lost in the tender reverie of my homeland.
     准备好了吗？开始提问吧！


    '''
    print(banner)
                                            
                                            

def start_robot():
    print_banner()
    print('>>', end='', flush=True)
    while True:
        try:
            msg = sys.stdin.readline().strip()
        except KeyboardInterrupt:
            print('\n再见!')
            break
        if not msg:
            print('>>', end='', flush=True)
            continue
        print('>>', end='', flush=True)
        # 获取并流式输出回复
        get_reply(msg)
        # 下一轮对话提示
        print('\n>>', end='', flush=True)

if __name__ == '__main__':
    start_robot()