from flask import Flask, render_template, request, jsonify, session
import base64
import io
import os
import json
from PIL import Image
from service import Service

app = Flask(__name__)
app.secret_key = os.urandom(24)  # 用于session加密

# 初始化服务
service = Service()

# 确保目录存在
if not os.path.exists('static'):
    os.makedirs('static')
if not os.path.exists('static/uploads'):
    os.makedirs('static/uploads')

@app.route('/')
def index():
    """渲染主页"""
    return render_template('index.html')

@app.route('/api/chat', methods=['POST'])
def chat():
    """处理聊天请求"""
    data = request.json
    prompt = data.get('message', '')
    image_data = data.get('image', None)
    history = data.get('history', [])
    
    # 确保历史记录包含系统提示
    has_system_prompt = False
    for msg in history:
        if msg.get('role') == 'system':
            has_system_prompt = True
            break
    
    if not has_system_prompt:
        # 添加默认的系统提示
        history.insert(0, {
            "role": "system",
            "content": "你是一个AI度量专家助手。你可以分析文本和图像的内容。能根据用户的需求，给出度量建议和洞察"
        })
    
    try:
        if image_data:
            # 使用图像调用API
            response = service.chat_with_image(
                text_prompt=prompt,
                image_base64=image_data,
                history=history
            )
        else:
            # 添加当前用户消息
            current_history = history.copy()
            current_history.append({"role": "user", "content": prompt})
            # 纯文本请求
            response = service.get_response(current_history)
            
        return jsonify({
            'status': 'success',
            'response': response
        })
    except Exception as e:
        return jsonify({
            'status': 'error',
            'message': str(e)
        }), 500

@app.route('/api/upload_image', methods=['POST'])
def upload_image():
    """处理图片上传"""
    if 'file' not in request.files:
        return jsonify({'status': 'error', 'message': '没有文件'})
    
    file = request.files['file']
    if file.filename == '':
        return jsonify({'status': 'error', 'message': '没有选择文件'})
    
    if file:
        filename = f"upload_{os.urandom(8).hex()}.png"
        filepath = os.path.join('static/uploads', filename)
        
        # 保存文件
        file.save(filepath)
        
        # 读取文件并转为base64
        with open(filepath, "rb") as img_file:
            img_str = base64.b64encode(img_file.read()).decode('utf-8')
        
        image_data = f"data:image/png;base64,{img_str}"
        
        return jsonify({
            'status': 'success',
            'image_data': image_data,
            'image_url': f"/static/uploads/{filename}"
        })

@app.route('/api/paste_image', methods=['POST'])
def paste_image():
    """处理粘贴的图片"""
    data = request.json
    image_data = data.get('image_data')
    
    if not image_data or not image_data.startswith('data:image'):
        return jsonify({'status': 'error', 'message': '无效的图片数据'})
    
    try:
        # 从base64解码
        image_data_parts = image_data.split(',')
        if len(image_data_parts) != 2:
            return jsonify({'status': 'error', 'message': '图片格式错误'})
        
        # 保存图片到文件
        filename = f"paste_{os.urandom(8).hex()}.png"
        filepath = os.path.join('static/uploads', filename)
        
        image_bytes = base64.b64decode(image_data_parts[1])
        with open(filepath, "wb") as f:
            f.write(image_bytes)
        
        return jsonify({
            'status': 'success',
            'image_data': image_data,
            'image_url': f"/static/uploads/{filename}"
        })
    except Exception as e:
        return jsonify({'status': 'error', 'message': str(e)})

if __name__ == '__main__':
    app.run(debug=True, port=5000)