#!/usr/bin/env python3
"""
语音转文字 - FunASR 版本（阿里达摩院）
支持中文识别、自动标点
"""
import os
import subprocess
import warnings

warnings.filterwarnings('ignore')

UPLOAD_FOLDER = '/home/ubuntu/nanocode/audio-recorder/uploads'
MODEL_CACHE = os.path.expanduser('~/.cache/modelscope')

class FunASRRecognizer:
    def __init__(self):
        self.model = None
        self.model_type = 'funasr' if self._check_model_cached() else 'vosk'
        
    def _check_model_cached(self):
        """检查 FunASR 模型是否已下载"""
        # 检查常用的模型路径
        possible_paths = [
            os.path.join(MODEL_CACHE, 'hub', 'damo', 'speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch'),
            os.path.join(MODEL_CACHE, 'hub', 'iic', 'speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch'),
        ]
        return any(os.path.exists(p) for p in possible_paths)
    
    def _load_model(self):
        """加载 FunASR 模型"""
        if self.model is None:
            try:
                from funasr import AutoModel
                print("🔄 加载 FunASR 模型...")
                self.model = AutoModel(
                    model="paraformer-zh",
                    vad_model="fsmn-vad",
                    punc_model="ct-punc",
                )
                print("✅ FunASR 模型加载完成")
                self.model_type = 'funasr'
            except Exception as e:
                print(f"⚠️ FunASR 加载失败: {e}, 回退到 Vosk")
                self.model_type = 'vosk'
                raise
        return self.model
    
    def convert_to_wav(self, input_path):
        """转换为 WAV 格式"""
        output_path = input_path + '.wav'
        try:
            cmd = [
                'ffmpeg', '-y', '-i', input_path,
                '-ar', '16000', '-ac', '1', '-f', 'wav',
                output_path
            ]
            result = subprocess.run(cmd, capture_output=True, timeout=30)
            if result.returncode == 0 and os.path.exists(output_path):
                return output_path
        except Exception as e:
            print(f"转换错误: {e}")
        return None
    
    def transcribe(self, audio_path):
        """转写音频"""
        # 先转换为 wav
        wav_path = self.convert_to_wav(audio_path)
        if not wav_path:
            return None, "音频转换失败"
        
        try:
            # 尝试使用 FunASR
            model = self._load_model()
            
            if self.model_type == 'funasr':
                # FunASR 识别
                res = model.generate(input=wav_path, batch_size_s=300)
                text = res[0]["text"]
            else:
                # 回退到 Vosk
                text, error = self._transcribe_vosk(wav_path)
                if error:
                    raise Exception(error)
            
            # 清理临时文件
            if os.path.exists(wav_path):
                os.remove(wav_path)
            
            return text, None
            
        except Exception as e:
            # 清理临时文件
            if os.path.exists(wav_path):
                os.remove(wav_path)
            return None, str(e)
    
    def _transcribe_vosk(self, wav_path):
        """备用：Vosk 识别"""
        try:
            from vosk import Model, KaldiRecognizer
            import json
            import wave
            
            model_path = '/home/ubuntu/nanocode/audio-recorder/models/cn-model'
            model = Model(model_path)
            
            wf = wave.open(wav_path, 'rb')
            rec = KaldiRecognizer(model, wf.getframerate())
            
            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 = json.loads(rec.FinalResult())
            if final.get('text'):
                results.append(final['text'])
            
            wf.close()
            return ' '.join(results), None
            
        except Exception as e:
            return None, f"Vosk 错误: {str(e)}"

# 全局识别器
_recognizer = None

def get_recognizer():
    global _recognizer
    if _recognizer is None:
        _recognizer = FunASRRecognizer()
    return _recognizer

def transcribe_audio(audio_path):
    """转写音频入口"""
    recognizer = get_recognizer()
    return recognizer.transcribe(audio_path)