#!/usr/bin/env python3.12

import soundfile as sf
import asyncio
import numpy as np
import time

from aces.connection_manager import ConnectionManager
from fastapi import WebSocket, WebSocketDisconnect


class AudioStreamer:
    def __init__(self, sound_file, ws_connection_mgr):
        self.sound_file = sound_file
        self.audio_file = None
        self.format_info = None
        self.chunk_size = 16384  # Larger chunk for better quality
        self.ws_connection_mgr = ws_connection_mgr
        self.broadcast_task = None
        self.init_audio_file()

    def init_audio_file(self):
        self.audio_file = sf.SoundFile(self.sound_file, 'rb')
        self.format_info = {
            'samplerate': self.audio_file.samplerate,
            'channels': self.audio_file.channels,
            'format': 'PCM',
            'subtype': str(self.audio_file.subtype)
        }

    async def broadcast_audio(self):
        try:
            chunk_duration = self.chunk_size / (self.audio_file.samplerate * self.audio_file.channels)
            target_interval = chunk_duration / 2  # Send chunks at twice playback rate for smooth buffering
            
            while True:
                if not self.ws_connection_mgr.active_connections:
                    await asyncio.sleep(0.1)
                    continue

                start_time = asyncio.get_event_loop().time()
                
                chunk = self.audio_file.read(self.chunk_size, dtype='float32')
                if len(chunk) == 0:
                    self.audio_file.seek(0)
                    continue

                await self.ws_connection_mgr.broadcast_raw(chunk.tobytes())
                
                # Calculate how long processing took and adjust sleep time
                elapsed = asyncio.get_event_loop().time() - start_time
                sleep_time = max(0, target_interval - elapsed)
                await asyncio.sleep(sleep_time)

        except Exception as e:
            print(f"Broadcast error: {e}")

    async def handle_client(self, ws: WebSocket):
        try:
            await self.ws_connection_mgr.connect(ws)
            await ws.send_json(self.format_info)

            if not self.broadcast_task or self.broadcast_task.done():
                self.broadcast_task = asyncio.create_task(self.broadcast_audio())

            while True:
                try:
                    await ws.receive_text()
                except WebSocketDisconnect:
                    break

        except Exception as e:
            print(f"Error in handle_client: {e}")
        finally:
            await self.ws_connection_mgr.disconnect(ws)