from flask import Flask, render_template, request, Response
from openai import OpenAI
from dotenv import load_dotenv
import os
import json
from datetime import datetime

load_dotenv()

app = Flask(__name__)
client = OpenAI(
    api_key=os.getenv('OPENAI_API_KEY'),
    base_url=os.getenv('OPENAI_BASE_URL')
)

class ChatBot:
    def __init__(self):
        self.messages = []
        # 加载系统提示词
        self.system_prompt = self.load_system_prompt()
        
    def load_system_prompt(self):
        # 从文件加载系统提示词
        try:
            with open('system_prompt.txt', 'r', encoding='utf-8') as file:
                return file.read()
        except Exception as e:
            print(f"Error loading system prompt: {e}")
            return ""
    
    def format_birth_info(self, birth_info):
        """格式化出生信息用于分析"""
        try:
            birth_date = datetime.strptime(birth_info['date'], '%Y年%m月%d日')
            formatted_date = birth_date.strftime('%Y年%m月%d日')
            return f"{formatted_date} {birth_info['time']}"
        except Exception as e:
            print(f"Error formatting birth info: {e}")
            return birth_info['raw']
    
    def prepare_analysis_prompt(self, user_data):
        """准备用于分析的完整提示词"""
        birth_info = self.format_birth_info(user_data['birth_info'])
        
        analysis_prompt = f"""根据以下信息进行命理分析：
出生信息：{birth_info}
咨询问题：{user_data['question']}

请基于此信息，结合八字、紫微斗数等进行分析。"""
        
        return analysis_prompt

    def get_stream_response(self, user_input):
        """处理用户输入并生成流式响应"""
        try:
            # 解析用户输入的JSON数据
            user_data = json.loads(user_input)
            
            # 如果是新的对话，添加系统提示词
            if not self.messages:
                self.messages.append({"role": "system", "content": self.system_prompt})
            
            # 准备分析提示词
            analysis_prompt = self.prepare_analysis_prompt(user_data)
            self.messages.append({"role": "user", "content": analysis_prompt})

            # 创建流式响应
            response = client.chat.completions.create(
                model=os.getenv('OPENAI_MODEL'),
                messages=self.messages,
                stream=True,
                temperature=0.7,
                max_tokens=2000
            )

            for chunk in response:
                if chunk.choices[0].delta.content:
                    yield f"data: {json.dumps({'content': chunk.choices[0].delta.content})}\n\n"

            # 保存完整对话到历史记录
            full_response = ''.join(chunk.choices[0].delta.content or '' 
                                  for chunk in response)
            self.messages.append({"role": "assistant", "content": full_response})

        except Exception as e:
            yield f"data: {json.dumps({'error': str(e)})}\n\n"

chatbot = ChatBot()

@app.route('/')
def home():
    return render_template('index.html')

@app.route('/chat', methods=['POST'])
def chat():
    try:
        user_message = request.get_json()
        return Response(
            chatbot.get_stream_response(json.dumps(user_message)),
            mimetype='text/event-stream'
        )
    except Exception as e:
        return jsonify({'error': str(e)}), 500

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860, debug=False)