#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""小狸v1.2 Web聊天服务器 - 当前训练中的模型"""
import sys, os, json, time
sys.path.insert(0, '/workspace/xiaoli-llm-chat')
import torch
from flask import Flask, request, jsonify, send_from_directory
from llm_chat_improved import LLMModel, LLMChat, build_vocabulary

print("=" * 50)
print("  小狸 v1.2 Web 服务加载中...")
print("=" * 50)

device = 'cpu'
model_dir = '/workspace/xiaoli-llm-chat'

# 构建当前vocab
with open(f'{model_dir}/xiaoli_merged_data_v2.json', 'r', encoding='utf-8') as f:
    all_data = json.load(f)
pairs = []
for d in all_data:
    q = d.get('question','') or d.get('input','') or d.get('prompt','') or ''
    a = d.get('answer','') or d.get('output','') or d.get('response','') or ''
    if q and a:
        pairs.append((q, a))
vocab = build_vocabulary(pairs)
print(f"词汇表: {len(vocab)} | 数据: {len(all_data)}条")

# 加载模型 (epoch 50 = 干净epoch 41 + vocab扩展)
ckpt_path = f'{model_dir}/llm_model_v2_resume_epoch_50.pth'
print(f"加载: {ckpt_path}")
ckpt = torch.load(ckpt_path, map_location=device)
old_vocab_size = ckpt.get('model_config', {}).get('vocab_size', 0)
cfg = ckpt['model_config']
print(f"断点vocab: {old_vocab_size}")

model = LLMModel(
    vocab_size=len(vocab),
    d_model=cfg['d_model'], n_heads=cfg['n_heads'],
    n_layers=cfg['n_layers'], d_ff=cfg['d_ff'],
    max_len=cfg['max_len'], dropout=cfg.get('dropout', 0.1),
    use_memory=False, use_thinking=False, use_moe=False,
    num_experts=1, top_k=1,
)

# 迁移权重
old_sd = ckpt['model_state_dict']
new_sd = model.state_dict()
if old_vocab_size != len(vocab):
    for key in new_sd:
        if key in old_sd:
            w = old_sd[key]
            if w.shape == new_sd[key].shape:
                new_sd[key] = w.clone()
            elif len(w.shape) == 2 and w.shape[0] == old_vocab_size:
                new_sd[key][:old_vocab_size] = w
            elif 'freq' in key:
                pass
            else:
                try: new_sd[key] = w.clone()
                except: pass
    model.load_state_dict(new_sd)
else:
    model.load_state_dict(old_sd)

model.to(device)
model.eval()
chat = LLMChat(model, vocab, device, max_length=cfg['max_len'])
print(f"✅ 模型加载完成！ | voc={len(vocab)} | max_len={cfg['max_len']}")

app = Flask(__name__)

@app.route('/')
def index():
    return send_from_directory(model_dir, 'chat.html')

@app.route('/api/chat', methods=['POST'])
def api_chat():
    data = request.get_json()
    question = data.get('question', '').strip()
    temp = float(data.get('temperature', 0.2))
    max_len = int(data.get('max_length', 100))
    history = data.get('history', None)
    
    if not question:
        return jsonify({'error': '问题不能为空'})
    
    # 如果有history，转换为对话格式
    history_pairs = None
    if history and isinstance(history, list) and len(history) > 0:
        history_pairs = [(h['q'], h['a']) for h in history]
    
    try:
        t0 = time.time()
        answer = chat.generate_response(question, max_length=max_len, temperature=temp, history=history_pairs)
        elapsed = time.time() - t0
        return jsonify({
            'question': question,
            'answer': answer,
            'time': f'{elapsed:.1f}s',
            'temperature': temp,
        })
    except Exception as e:
        return jsonify({'error': str(e)})

@app.route('/api/info')
def model_info():
    return jsonify({
        'name': '小狸 v1.2',
        'version': 'epoch 52 (架构扩展: 7.9M→21.3M, max_len=256)',
        'vocab_size': len(vocab),
        'max_length': cfg['max_len'],
        'd_model': cfg['d_model'],
        'layers': cfg['n_layers'],
        'heads': cfg['n_heads'],
        'd_ff': cfg['d_ff'],
        'data_size': f'{len(all_data):,}条',
        'training': '新架构训练中 epoch 1→200',
        'features': ['多轮对话(1,694条)', '长文本(758条)', '我不知道(1,140条)'],
    })

if __name__ == '__main__':
    print("\n" + "=" * 50)
    print("  小狸 v1.2 聊天服务已启动！")
    print("  访问地址: http://0.0.0.0:5000")
    print("  建议温度: 0.1 ~ 0.2")
    print("=" * 50)
    app.run(host='0.0.0.0', port=5000, debug=False, threaded=True)