import pyaudio
import wave
import threading
import time
import numpy as np
from xinference.client import Client

class StreamingASR:
    def __init__(self, server_url="http://124.93.18.180:9997", model_name="whisper-large-v3"):
        self.client = Client(server_url)
        self.model = self.client.get_model(model_name)
        self.is_recording = False
        self.audio_data = []
        self.chunk_duration = 5
        self.sample_rate = 16000
        self.channels = 1
        self.chunk_size = 1024
        self.silence_threshold = 500
        self.silence_duration = 1.0
        self.silence_counter = 0
        
    def calculate_audio_energy(self, audio_chunk):
        audio_array = np.frombuffer(audio_chunk, dtype=np.int16)
        energy = np.sqrt(np.mean(audio_array.astype(np.float32) ** 2))
        return energy
    
    def record_audio(self):
        audio = pyaudio.PyAudio()
        stream = audio.open(
            format=pyaudio.paInt16,
            channels=self.channels,
            rate=self.sample_rate,
            input=True,
            frames_per_buffer=self.chunk_size
        )
        
        print("开始录音... (检测到1秒静音自动结束)")
        self.is_recording = True
        self.silence_counter = 0
        
        try:
            while self.is_recording:
                data = stream.read(self.chunk_size)
                self.audio_data.append(data)
                
                energy = self.calculate_audio_energy(data)
                if energy < self.silence_threshold:
                    self.silence_counter += self.chunk_size / self.sample_rate
                else:
                    self.silence_counter = 0
                
                if self.silence_counter >= self.silence_duration:
                    print(f"检测到{self.silence_duration}秒静音，结束录音")
                    self.is_recording = False
                    break
                
                if len(self.audio_data) * self.chunk_size >= self.sample_rate * 2 * self.chunk_duration:
                    self.process_audio_chunk()
                    
        except KeyboardInterrupt:
            print("\n停止录音")
        finally:
            stream.stop_stream()
            stream.close()
            audio.terminate()
            
            if self.audio_data:
                self.process_audio_chunk()
    
    def process_audio_chunk(self):
        import io
        audio_bytes = b''.join(self.audio_data)
        
        audio_file = io.BytesIO(audio_bytes)
        audio_file.name = "audio.wav"
        
        try:
            transcription = self.model.transcriptions(audio_file.read())
            if transcription.text.strip():
                print(f"识别结果: {transcription.text}")
        except Exception as e:
            print(f"识别错误: {e}")
        
        self.audio_data = []
    
    def start_streaming(self):
        record_thread = threading.Thread(target=self.record_audio)
        record_thread.start()
        return record_thread

if __name__ == "__main__":
    asr = StreamingASR()
    thread = asr.start_streaming()
    thread.join()
