import sounddevice as sd
import soundfile as sf
import numpy as np
from queue import Queue
import threading
import time

import whisper

# 启动加载 whisper 语言训练模型
model = whisper.load_model("medium")

def record_audio(stop_event, data_queue):
    '''
    使用 sounddevice 库录制音频
    '''
    def callback(indata, frames, time, status):
        '''
        sounddevice.RawInputStream 回调函数

        indata 是当前缓冲区的输入数据
        frames 指示当前缓冲区中有多少帧数据
        time 提供时间信息，例如当前回调的时间戳。
        status 指示回调执行过程中的状态或错误信息
        '''

        # 如果 status 非空，即存在错误或警告，它会被打印到控制台
        if status:
            print(status)
        data_queue.put(bytes(indata))

    # 采样率16000，数据类型为16位整数，通道位1的单声道录音
    with sounddevice.RawInputStream(samplerate=16000, dtype="int16", channels=1, callback=callback):
        while not stop_event.is_set():
            # 减轻 CPU 负担，每隔 0.1 秒检查 stop_event 状态
            time.sleep(0.1)

if __name__ == '__main__':
    print('Voice service started (Press Ctrl+C to exit)')

    try:
        while True:
            print('Press Enter to start recording, again to stop')

            # --------------- 启动音频录制 ---------------

            # 创建队列存储音频字节，启动线程开始接受音频输入
            data_queue = Queue()
            stop_event = threading.Event()
            recording_thread = threading.Thread(
                target=record_audio,
                args=(stop_event, data_queue),
            )
            recording_thread.start()

            # 使用 input() 函数等待用户输入，暂停程序执行，知道用户按下回车键
            input()
            # 设置线程旗标为 True，停止录制过程
            stop_event.set()
            # 等待recording_thread线程完成，即record_audio函数完成其执行
            # 在程序结束前的一个整理步骤，确保所有资源得以适当释放且没有遗留工作
            recording_thread.join()

            # --------------- 音频数据处理 ---------------

            audio_data = b"".join(list(data_queue.queue))
            # 用 numpy 库来将字节缓冲区（audio_data）转换为一个 NumPy 数组，并做类型转换和标准化处理
            audio_np = (
                np.frombuffer(audio_data, dtype=np.int16).astype(np.float32) / 32768.0
            )

            if audio_np.size > 0:
                # 使用 openai whisper 识别音频字节数据，并输出到控制台
                result = model.transcribe(audio_np, fp16=False)
                print(result["text"])

                sd.wait()
            else:
                print("No audio recorded. Please ensure your microphone is working.")
    except KeyboardInterrupt:
        print("\nVoice service ended")

# --- snip --- 
# 采样率44100
    with sounddevice.RawInputStream(samplerate=44100, dtype="int16", channels=1, callback=callback):
# --- snip --- 

if audio_np.size > 0:
    # 播放录制的音频
    sd.play(audio_np, 44100)
    sd.wait()

import soundfile as sf
# --- snip --- 

if audio_np.size > 0:
    # 写入录制的音频
    sf.write('output.wav', audio_np, 44100, format='WAV')