import socket
import wave
import pyaudio
import threading
import time
import struct

class AudioReceiver:
    def __init__(self, host='0.0.0.0', port=12345):
        # Audio Configuration
        self.SAMPLE_RATE = 16000
        self.CHANNELS = 1
        self.SAMPLE_WIDTH = 2  # 16-bit = 2 bytes
        self.CHUNK_SIZE = 1024
        
        self.host = host
        self.port = port
        self.running = False
        
        # Initialize PyAudio
        self.p = pyaudio.PyAudio()
        self.stream = self.p.open(
            format=pyaudio.paInt16,
            channels=self.CHANNELS,
            rate=self.SAMPLE_RATE,
            output=True,
            frames_per_buffer=self.CHUNK_SIZE
        )
        
        # Buffer for audio data
        self.audio_buffer = bytearray()
        self.buffer_lock = threading.Lock()
        
    def process_audio_data(self, data):
        with self.buffer_lock:
            self.audio_buffer.extend(data)
            
            # Process complete chunks
            while len(self.audio_buffer) >= self.CHUNK_SIZE:
                chunk = self.audio_buffer[:self.CHUNK_SIZE]
                self.stream.write(bytes(chunk))
                self.audio_buffer = self.audio_buffer[self.CHUNK_SIZE:]
    
    def start(self):
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.sock.bind((self.host, self.port))
        self.sock.settimeout(1.0)  # 1 second timeout
        
        self.running = True
        print(f"Listening on {self.host}:{self.port}")
        
        # Create WAV file for recording
        self.wave_file = wave.open('received_audio.wav', 'wb')
        self.wave_file.setnchannels(self.CHANNELS)
        self.wave_file.setsampwidth(self.SAMPLE_WIDTH)
        self.wave_file.setframerate(self.SAMPLE_RATE)
        
        while self.running:
            try:
                data, addr = self.sock.recvfrom(self.CHUNK_SIZE)
                self.process_audio_data(data)
                self.wave_file.writeframes(data)
                
            except socket.timeout:
                continue
            except Exception as e:
                print(f"Error: {e}")
                break
                
    def stop(self):
        self.running = False
        self.stream.stop_stream()
        self.stream.close()
        self.p.terminate()
        self.wave_file.close()
        self.sock.close()

def main():
    receiver = AudioReceiver()
    try:
        receiver.start()
    except KeyboardInterrupt:
        print("\nStopping receiver")
        receiver.stop()

if __name__ == "__main__":
    main()
