#!/usr/bin/env python3
"""
录音后端服务 - 集成语音转文字
"""
from flask import Flask, request, jsonify, send_from_directory
from flask_cors import CORS
from werkzeug.utils import secure_filename
import os
import ssl
import wave
import json
from datetime import datetime

app = Flask(__name__)
CORS(app)

# 配置
UPLOAD_FOLDER = '/home/ubuntu/nanocode/audio-recorder/uploads'
MODEL_PATH = '/home/ubuntu/nanocode/audio-recorder/models/cn-model'
ALLOWED_EXTENSIONS = {'webm', 'mp4', 'wav', 'mp3', 'ogg'}
MAX_CONTENT_LENGTH = 100 * 1024 * 1024

app.config['MAX_CONTENT_LENGTH'] = MAX_CONTENT_LENGTH
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

# 全局 ASR 模型（延迟加载）
funasr_model = None
vosk_model = None

def get_funasr_model():
    """加载 FunASR 模型（阿里，中文识别效果好）- 优先使用"""
    global funasr_model
    if funasr_model is None:
        try:
            from funasr import AutoModel
            print("🔄 加载 FunASR 模型（阿里达摩院，中文优先）...")
            funasr_model = AutoModel(
                model="paraformer-zh",
                vad_model="fsmn-vad",
                punc_model="ct-punc",
            )
            print("✅ FunASR 模型加载完成")
        except Exception as e:
            print(f"⚠️ FunASR 加载失败: {e}")
            return None
    return funasr_model

def get_vosk_model():
    """加载 Vosk 模型（备用）"""
    global vosk_model
    if vosk_model is None:
        try:
            from vosk import Model
            if os.path.exists(MODEL_PATH):
                print(f"🔄 加载 Vosk 模型（备用）...")
                vosk_model = Model(MODEL_PATH)
                print("✅ Vosk 模型加载完成")
            else:
                print(f"⚠️ Vosk 模型不存在: {MODEL_PATH}")
                return None
        except Exception as e:
            print(f"⚠️ Vosk 加载失败: {e}")
            return None
    return vosk_model

def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

def convert_to_wav(input_path, output_path):
    """将音频转换为 WAV 格式（16kHz, 单声道）"""
    import subprocess
    try:
        cmd = [
            'ffmpeg', '-y', '-i', input_path,
            '-ar', '16000', '-ac', '1', '-f', 'wav',
            output_path
        ]
        subprocess.run(cmd, check=True, capture_output=True, timeout=30)
        return True
    except Exception as e:
        print(f"转换失败: {e}")
        return False

def transcribe_with_funasr(audio_path):
    """使用 FunASR 转写（优先，中文效果好）"""
    model = get_funasr_model()
    if not model:
        return None, "FunASR 模型未加载"
    
    try:
        res = model.generate(input=audio_path, batch_size_s=300)
        text = res[0]["text"] if res and len(res) > 0 else ""
        return text, None
    except Exception as e:
        return None, f"FunASR 错误: {str(e)}"

def transcribe_with_vosk(audio_path):
    """使用 Vosk 转写音频"""
    from vosk import KaldiRecognizer
    
    model = get_asr_model()
    if not model:
        return None, "模型未加载"
    
    try:
        # 转换为 WAV
        wav_path = audio_path + '.wav'
        if not convert_to_wav(audio_path, wav_path):
            return None, "音频转换失败"
        
        # 读取 WAV 文件
        wf = wave.open(wav_path, 'rb')
        if wf.getnchannels() != 1 or wf.getsampwidth() != 2 or wf.getcomptype() != 'NONE':
            return None, "音频格式不支持"
        
        # 创建识别器
        rec = KaldiRecognizer(model, wf.getframerate())
        rec.SetWords(True)
        
        # 识别
        results = []
        while True:
            data = wf.readframes(4000)
            if len(data) == 0:
                break
            if rec.AcceptWaveform(data):
                result = json.loads(rec.Result())
                if result.get('text'):
                    results.append(result['text'])
        
        # 获取最终结果
        final_result = json.loads(rec.FinalResult())
        if final_result.get('text'):
            results.append(final_result['text'])
        
        # 清理临时文件
        wf.close()
        if os.path.exists(wav_path):
            os.remove(wav_path)
        
        return ' '.join(results), None
        
    except Exception as e:
        return None, str(e)

@app.route('/api/upload', methods=['POST'])
def upload_audio():
    """接收音频文件上传"""
    if 'audio' not in request.files:
        return jsonify({'error': '没有文件'}), 400
    
    file = request.files['audio']
    if file.filename == '':
        return jsonify({'error': '文件名为空'}), 400
    
    if file and allowed_file(file.filename):
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        filename = f"recording_{timestamp}.webm"
        filepath = os.path.join(UPLOAD_FOLDER, filename)
        
        file.save(filepath)
        file_size = os.path.getsize(filepath)
        
        return jsonify({
            'success': True,
            'filename': filename,
            'filepath': filepath,
            'size': file_size,
            'size_human': f"{file_size / 1024:.1f} KB",
            'timestamp': timestamp
        }), 200
    
    return jsonify({'error': '不支持的文件类型'}), 400

@app.route('/api/transcribe', methods=['POST'])
def transcribe():
    """语音转文字 - 优先使用 FunASR（中文效果好），失败回退 Vosk"""
    if 'audio' not in request.files:
        return jsonify({'error': '没有音频文件'}), 400
    
    file = request.files['audio']
    if file.filename == '':
        return jsonify({'error': '文件名为空'}), 400
    
    try:
        # 保存上传的音频
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        filename = f"recording_{timestamp}.webm"
        filepath = os.path.join(UPLOAD_FOLDER, filename)
        file.save(filepath)
        
        # 转换为 WAV 格式（两种引擎都需要）
        wav_path = filepath + '.wav'
        if not convert_to_wav(filepath, wav_path):
            return jsonify({
                'success': False,
                'error': '音频格式转换失败'
            }), 500
        
        # 1. 优先使用 FunASR（阿里，中文识别效果好）
        text, error = transcribe_with_funasr(wav_path)
        engine = "funasr"
        
        # 2. FunASR 失败，回退到 Vosk
        if not text:
            print(f"⚠️ FunASR 失败({error})，回退到 Vosk...")
            text, error = transcribe_with_vosk(wav_path)
            engine = "vosk"
        
        # 清理临时 WAV 文件
        if os.path.exists(wav_path):
            os.remove(wav_path)
        
        if error and not text:
            return jsonify({
                'success': False,
                'error': error,
                'filename': filename
            }), 500
        
        return jsonify({
            'success': True,
            'text': text or '（未识别到语音）',
            'filename': filename,
            'engine': engine
        })
        
    except Exception as e:
        return jsonify({
            'success': False,
            'error': str(e)
        }), 500

@app.route('/api/recordings', methods=['GET'])
def list_recordings():
    """列出所有录音文件"""
    recordings = []
    
    for filename in sorted(os.listdir(UPLOAD_FOLDER), reverse=True):
        if allowed_file(filename):
            filepath = os.path.join(UPLOAD_FOLDER, filename)
            stat = os.stat(filepath)
            recordings.append({
                'filename': filename,
                'size': stat.st_size,
                'size_human': f"{stat.st_size / 1024:.1f} KB",
                'created': datetime.fromtimestamp(stat.st_mtime).isoformat()
            })
    
    return jsonify({
        'count': len(recordings),
        'recordings': recordings
    })

@app.route('/api/recordings/<filename>', methods=['GET'])
def get_recording(filename):
    """获取单个录音文件"""
    return send_from_directory(UPLOAD_FOLDER, filename)

@app.route('/api/health', methods=['GET'])
def health_check():
    """健康检查"""
    model_exists = os.path.exists(MODEL_PATH)
    return jsonify({
        'status': 'ok',
        'model_loaded': model_exists,
        'model_path': MODEL_PATH,
        'upload_folder': UPLOAD_FOLDER,
        'recordings_count': len([f for f in os.listdir(UPLOAD_FOLDER) if allowed_file(f)])
    })

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

@app.route('/<path:path>')
def static_files(path):
    return send_from_directory('../frontend', path)

if __name__ == '__main__':
    context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
    context.load_cert_chain(
        '/home/ubuntu/nanocode/audio-recorder/cert.pem',
        '/home/ubuntu/nanocode/audio-recorder/key.pem'
    )
    
    print(f"🎙️ HTTPS 录音服务启动中...")
    print(f"📁 上传目录: {UPLOAD_FOLDER}")
    print(f"🔤 模型路径: {MODEL_PATH}")
    print(f"🔒 HTTPS 地址: https://129.211.173.87:8891")
    
    if not os.path.exists(MODEL_PATH):
        print(f"⚠️ 警告: ASR 模型未找到，请先下载模型")
    
    app.run(
        host='0.0.0.0',
        port=8891,
        ssl_context=context,
        debug=False
    )