import asyncio
import websockets
import argparse
import json
import logging
from funasr import AutoModel
import os

# 设置日志级别
logger = logging.getLogger(__name__)
logger.setLevel(logging.CRITICAL)

# 解析命令行参数
parser = argparse.ArgumentParser()
parser.add_argument("--host", type=str, default="0.0.0.0", help="host ip, localhost, 0.0.0.0")
parser.add_argument("--port", type=int, default=10197, help="grpc server port")
parser.add_argument("--ngpu", type=int, default=1, help="0 for cpu, 1 for gpu")
parser.add_argument("--gpu_id", type=int, default=0, help="specify which gpu device to use")
args = parser.parse_args()

# 初始化模型
print("model loading")
asr_model = AutoModel(model="paraformer-zh", model_revision="v2.0.4",
                      vad_model="fsmn-vad", vad_model_revision="v2.0.4",
                      punc_model="ct-punc-c", punc_model_revision="v2.0.4",
                      device=f"cuda:{args.gpu_id}" if args.ngpu else "cpu", disable_update=True)
                    #   ,disable_update=True
print("model loaded")
websocket_users = {}
task_queue = asyncio.Queue()

async def ws_serve(websocket, path):
    global websocket_users
    user_id = id(websocket)
    websocket_users[user_id] = websocket
    try:
        async for message in websocket:
            if isinstance(message, str):
                data = json.loads(message)
                if 'url' in data:
                    # 处理文件URL
                    await task_queue.put((websocket, data['url'], 'url'))
                elif 'audio_data' in data:
                    # 处理音频数据
                    await task_queue.put((websocket, data, 'audio_data'))
    except websockets.exceptions.ConnectionClosed as e:
        logger.info(f"Connection closed: {e.reason}")
    except Exception as e:
        logger.error(f"Unexpected error: {e}")
    finally:
        logger.info(f"Cleaning up connection for user {user_id}")
        if user_id in websocket_users:
            del websocket_users[user_id]
        await websocket.close()
        logger.info("WebSocket closed")

async def worker():
    while True:
        task_data = await task_queue.get()
        websocket = task_data[0]
        
        if websocket.open:
            if len(task_data) == 3:  # 新格式: (websocket, data, type)
                data, data_type = task_data[1], task_data[2]
                if data_type == 'url':
                    await process_wav_file(websocket, data)
                elif data_type == 'audio_data':
                    await process_audio_data(websocket, data)
            else:  # 兼容旧格式: (websocket, url)
                await process_wav_file(websocket, task_data[1])
        else:
            logger.info("WebSocket connection is already closed when trying to process file")
        task_queue.task_done()

async def process_wav_file(websocket, url):
    # 热词
    param_dict = {"sentence_timestamp": False}
    with open("data/hotword.txt", "r", encoding="utf-8") as f:
        lines = f.readlines()
        lines = [line.strip() for line in lines]
    hotword = " ".join(lines)
    print(f"热词：{hotword}")
    param_dict["hotword"] = hotword
    wav_path = url
    try:
        res = asr_model.generate(input=wav_path, is_final=True, **param_dict)
        if res:
            if 'text' in res[0] and websocket.open:
                await websocket.send(res[0]['text'])
    except Exception as e:
        print(f"Error during model.generate: {e}")
    finally:
        if os.path.exists(wav_path):
            os.remove(wav_path)

async def process_audio_data(websocket, data):
    """处理音频数据"""
    import base64
    import tempfile
    
    try:
        # 获取音频数据
        audio_data = data.get('audio_data')
        filename = data.get('filename', 'audio.wav')
        
        if not audio_data:
            await websocket.send(json.dumps({"error": "No audio data provided"}))
            return
        
        # 解码Base64音频数据
        audio_bytes = base64.b64decode(audio_data)
        
        # 创建临时文件
        with tempfile.NamedTemporaryFile(delete=False, suffix='.wav') as temp_file:
            temp_file.write(audio_bytes)
            temp_path = temp_file.name
        
        print(f"处理音频文件: {filename}, 临时路径: {temp_path}")
        
        # 热词配置
        param_dict = {"sentence_timestamp": False}
        try:
            with open("data/hotword.txt", "r", encoding="utf-8") as f:
                lines = f.readlines()
                lines = [line.strip() for line in lines]
            hotword = " ".join(lines)
            print(f"热词：{hotword}")
            param_dict["hotword"] = hotword
        except FileNotFoundError:
            print("热词文件不存在，跳过热词配置")
        
        # 进行语音识别
        res = asr_model.generate(input=temp_path, is_final=True, **param_dict)
        if res and websocket.open:
            if 'text' in res[0]:
                result_text = res[0]['text']
                print(f"识别结果: {result_text}")
                await websocket.send(result_text)
            else:
                await websocket.send("识别失败：无法获取文本结果")
        
    except Exception as e:
        print(f"处理音频数据时出错: {e}")
        if websocket.open:
            await websocket.send(f"识别错误: {str(e)}")
    finally:
        # 清理临时文件
        if 'temp_path' in locals() and os.path.exists(temp_path):
            os.remove(temp_path)

async def main():
    server = await websockets.serve(ws_serve, args.host, args.port, ping_interval=10)
    worker_task = asyncio.create_task(worker())
    
    try:
        # 保持服务器运行，直到被手动中断
        print(f"ASR服务器已启动，监听地址: {args.host}:{args.port}")
        await asyncio.Future()  # 永久等待，直到程序被中断
    except asyncio.CancelledError:
        print("服务器正在关闭...")
    finally:
        # 清理资源
        worker_task.cancel()
        try:
            await worker_task
        except asyncio.CancelledError:
            pass
        server.close()
        await server.wait_closed()

# 使用 asyncio 运行主函数
asyncio.run(main())