#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
小狸聊天 Web 服务 v2 - 支持多轮对话
"""
import sys, os, json, re
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from flask import Flask, render_template, request, jsonify
import torch
from llm_chat_improved import (
    CharTokenizer, build_vocabulary, LLMModel, LLMChat, load_training_data
)

app = Flask(__name__)

# 全局模型
llm_chat = None
vocab = None

def clean_text(text):
    """清洗文本"""
    text = text.replace('\n', '')
    text = re.sub(r'\*\*([^*]+)\*\*', r'\1', text)
    text = re.sub(r'\([^)]*[ω｡ﾉ♡^][^)]*\)', '', text)
    emoji_pattern = re.compile(
        '[\U0001F600-\U0001F64F\U0001F300-\U0001F5FF\U0001F680-\U0001F6FF\U0001F1E0-\U0001F1FF\U00002702-\U000027B0\U0001f926-\U0001f937\u2640-\u2642\u2600-\u2B55\u200d\u23cf\u23e9\u231a\ufe0f\u3030\u2728\u2714\u2716\u27a1\u2b05\u2b06\u2b07\u2b1b\u2b1c\u2b50\u2b55\u2934\u2935\u25aa\u25ab\u25b6\u25c0\u25fb\u25fc\u25fd\u25fe]+',
        flags=re.UNICODE
    )
    text = emoji_pattern.sub('', text)
    text = re.sub(r'[ \t]+', ' ', text)
    return text.strip()

def load_model():
    """加载训练好的模型"""
    global llm_chat, vocab
    
    print("🐱 正在加载小狸模型...")
    
    device = 'cpu'
    
    # 直接从模型文件加载 vocab 和配置 (v4大模型)
    checkpoint = torch.load('xiaoli_final_v4.pth', map_location=device)
    vocab = checkpoint['vocab']
    config = checkpoint.get('model_config', {})
    
    print(f"   词汇表大小: {len(vocab)}")
    
    # 创建模型 - 使用保存的配置（v4大模型带MoE+Thinking+Memory）
    model = LLMModel(
        vocab_size=config.get('vocab_size', len(vocab)),
        d_model=config.get('d_model', 384),
        n_heads=config.get('n_heads', 8),
        n_layers=config.get('n_layers', 6),
        d_ff=config.get('d_ff', 1536),
        max_len=config.get('max_len', 128),
        dropout=0.0,
        use_memory=config.get('use_memory', True),
        use_thinking=config.get('use_thinking', True),
        use_moe=config.get('use_moe', True),
        num_experts=config.get('num_experts', 4),
        top_k=config.get('top_k', 2),
    )
    
    # 加载权重
    model.load_state_dict(checkpoint['model_state_dict'])
    
    llm_chat = LLMChat(model, vocab, device, 128)
    print("✅ 小狸已就绪！")

# 存储会话历史
session_history = {}

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

@app.route('/chat', methods=['POST'])
def chat():
    data = request.json
    message = data.get('message', '').strip()
    session_id = data.get('session_id', 'default')
    
    if not message:
        return jsonify({'response': '请输入内容哦~'})
    
    if llm_chat is None:
        return jsonify({'response': '模型加载中，请稍后再试'})
    
    try:
        # 获取会话历史
        history = session_history.get(session_id, [])
        
        # 构建带上下文的输入（使用与训练一致的格式）
        # 训练格式: <SOS>问题<SEP><START>
        # 多轮格式: <SOS>历史问题1<SEP>历史回答1<SEP>历史问题2<SEP>历史回答2<SEP>当前问题<SEP><START>
        if history and len(history) >= 2:
            # 只保留最近2轮对话（4条记录）
            recent_history = history[-4:]
            # 提取问题和答案
            context_parts = []
            for i in range(0, len(recent_history), 2):
                if i + 1 < len(recent_history):
                    q = recent_history[i].replace("用户：", "")
                    a = recent_history[i+1].replace("小狸：", "")
                    context_parts.append(f"{q}<SEP>{a}")
            
            if context_parts:
                context = "<SEP>".join(context_parts)
                full_input = f"{context}<SEP>{message}"
            else:
                full_input = message
        else:
            full_input = message
        
        # 限制输入长度（字符数）
        if len(full_input) > 80:
            full_input = full_input[-80:]
        
        response = llm_chat.generate_response(full_input, max_length=50, temperature=0.3)
        
        # 更新历史
        history.append(f"用户：{message}")
        history.append(f"小狸：{response}")
        # 只保留最近6轮（12条记录）
        session_history[session_id] = history[-12:]
        
        return jsonify({'response': response})
    except Exception as e:
        import traceback
        return jsonify({'response': f'出错了: {str(e)}\n{traceback.format_exc()}'})

@app.route('/clear', methods=['POST'])
def clear_history():
    data = request.json
    session_id = data.get('session_id', 'default')
    session_history[session_id] = []
    return jsonify({'status': 'ok'})

if __name__ == '__main__':
    load_model()
    print("\n🌐 启动 Web 服务: http://0.0.0.0:5000")
    app.run(host='0.0.0.0', port=5000, debug=False)