#!/usr/bin/env python3
"""
Audio Output Module - Handles audio playback using pygame

This module provides audio output functionality for the PrimVoices debugger.
It uses prebuffering to ensure smooth, gapless audio playback.
"""

import threading
import asyncio
import json
import time
import numpy as np
from typing import Any, Dict, Optional, List

from vr_cli.utils.utils import print_error
from vr_cli.utils.config import OUTPUT_SAMPLE_RATE

try:
    import sys
    import os
    # Temporarily redirect stdout to suppress pygame welcome message
    with open(os.devnull, 'w') as devnull:
        old_stdout = sys.stdout
        sys.stdout = devnull
        import pygame
        sys.stdout = old_stdout
except ImportError as e:
    print(f"Missing required dependency: {e}")
    print("Please install required packages:")
    print("pip install pygame")
    raise


class AudioOutput:
    """Handles audio playback using pygame with prebuffering for smooth playback"""

    # Prebuffering settings (similar to JS WebSocketClient)
    MIN_PREBUFFER_SECONDS = 0.3  # Wait for this much audio before starting playback
    MAX_BUFFER_SECONDS = 2.0     # Maximum buffer size to prevent memory issues
    CHUNK_COMBINE_COUNT = 4      # Number of chunks to combine for smoother playback

    def __init__(self):
        self.pygame_initialized = False
        self.is_playing = False
        self.is_buffering = True  # Start in buffering mode
        self.playback_thread = None

        # Separate queues for audio and marks
        self.audio_buffer: List[np.ndarray] = []  # Raw audio chunks waiting to be played
        self.mark_queue: List[tuple] = []  # Mark events with timestamps

        # Lock for thread-safe queue access
        self.buffer_lock = threading.Lock()

        # Echo cancellation: store recent output audio for cancellation
        self.echo_buffer: List[np.ndarray] = []
        self.echo_buffer_max_size = 50  # Keep last 50 chunks (~1 second at 20ms chunks)

        # Track buffered duration
        self.buffered_samples = 0
        self.samples_played = 0

    def _ensure_pygame_initialized(self):
        """Initialize pygame mixer if not already initialized"""
        if not self.pygame_initialized:
            try:
                # Use a larger buffer for smoother playback
                pygame.mixer.init(frequency=OUTPUT_SAMPLE_RATE, size=-16, channels=1, buffer=2048)
                self.pygame_initialized = True
            except Exception as e:
                print_error(f"Failed to initialize pygame mixer: {e}")
                raise

    def _get_buffered_seconds(self) -> float:
        """Calculate total buffered audio duration in seconds"""
        with self.buffer_lock:
            total_samples = sum(len(chunk) for chunk in self.audio_buffer)
            return total_samples / OUTPUT_SAMPLE_RATE

    def play_pcm_bytes(
        self,
        audio_bytes: bytes,
        sample_rate: int = OUTPUT_SAMPLE_RATE
    ):
        """Queue raw 16-bit little-endian PCM bytes (mono) for playback."""
        # Ensure pygame mixer is initialized
        self._ensure_pygame_initialized()

        # Convert bytes to numpy array
        pcm = np.frombuffer(audio_bytes, dtype=np.int16).astype(np.float32) / 32767.0

        # Resample if needed
        if sample_rate != OUTPUT_SAMPLE_RATE:
            pcm = self._resample_audio(pcm, sample_rate, OUTPUT_SAMPLE_RATE)

        # Add to buffer
        with self.buffer_lock:
            self.audio_buffer.append(pcm)
            self.buffered_samples += len(pcm)

        # Check if we should start playback
        buffered_seconds = self._get_buffered_seconds()

        if self.is_buffering and buffered_seconds >= self.MIN_PREBUFFER_SECONDS:
            # We have enough audio buffered, start playback
            self.is_buffering = False
            if not self.is_playing:
                self._start_playback_thread()
        elif not self.is_playing and not self.is_buffering:
            # Playback stopped but we're not in buffering mode, restart
            self._start_playback_thread()

    def _start_playback_thread(self):
        """Start the audio playback thread"""
        if self.is_playing:
            return

        # Clear echo buffer when starting new playback
        self.echo_buffer.clear()
        self.samples_played = 0

        # Capture the event loop from the main thread for mark events
        try:
            loop = asyncio.get_event_loop()
        except RuntimeError:
            loop = None

        def _playback_worker():
            try:
                self.is_playing = True

                while self.is_playing:
                    # Get combined audio chunk from buffer
                    combined_chunk = self._get_combined_chunk()

                    if combined_chunk is None:
                        # No audio available, check if we should wait or stop
                        if self._get_buffered_seconds() < 0.05:
                            # Buffer is nearly empty, wait a bit for more data
                            time.sleep(0.02)
                            # Check again
                            if self._get_buffered_seconds() < 0.01 and len(self.audio_buffer) == 0:
                                # Still empty, stop playback
                                break
                        continue

                    # Process any pending mark events
                    self._process_marks(loop)

                    # Convert to int16 for pygame
                    pcm_int16 = (combined_chunk * 32767).astype(np.int16)

                    # Store in echo buffer
                    self.echo_buffer.append(combined_chunk)
                    if len(self.echo_buffer) > self.echo_buffer_max_size:
                        self.echo_buffer.pop(0)

                    # Create and play the sound
                    try:
                        sound = pygame.sndarray.make_sound(pcm_int16)
                        sound.play()

                        # Wait for playback to complete (with small buffer for gapless)
                        sound_duration = len(pcm_int16) / OUTPUT_SAMPLE_RATE
                        wait_time = sound_duration - 0.02  # Start next chunk slightly early
                        if wait_time > 0:
                            time.sleep(wait_time)

                        # Wait for mixer to be ready
                        while pygame.mixer.get_busy() and self.is_playing:
                            time.sleep(0.005)

                    except Exception as e:
                        print_error(f"Error playing audio chunk: {e}")
                        break

            except Exception as e:
                print_error(f"Error in audio playback worker: {e}")
            finally:
                self.is_playing = False
                self.is_buffering = True  # Reset to buffering mode for next playback

        self.playback_thread = threading.Thread(target=_playback_worker, daemon=True)
        self.playback_thread.start()

    def _get_combined_chunk(self) -> Optional[np.ndarray]:
        """Get multiple chunks combined into one for smoother playback"""
        with self.buffer_lock:
            if not self.audio_buffer:
                return None

            # Combine multiple small chunks into one larger chunk
            chunks_to_combine = min(self.CHUNK_COMBINE_COUNT, len(self.audio_buffer))
            if chunks_to_combine == 0:
                return None

            chunks = [self.audio_buffer.pop(0) for _ in range(chunks_to_combine)]
            combined = np.concatenate(chunks)
            self.samples_played += len(combined)
            return combined

    def _process_marks(self, loop):
        """Process any pending mark events"""
        with self.buffer_lock:
            marks_to_send = list(self.mark_queue)
            self.mark_queue.clear()

        for mark_event, websocket in marks_to_send:
            if loop is not None:
                try:
                    asyncio.run_coroutine_threadsafe(
                        websocket.send(json.dumps(mark_event)),
                        loop
                    )
                except Exception as e:
                    print_error(f"Error sending mark event: {e}")

    def add_mark_event(self, mark_event: Dict[str, Any], websocket: Any):
        """Add a mark event to be sent when audio reaches this point"""
        with self.buffer_lock:
            self.mark_queue.append((mark_event, websocket))

    def stop_playback(self):
        """Stop audio playback and clear the buffer"""
        self.is_playing = False
        self.is_buffering = True

        # Stop pygame mixer immediately
        if self.pygame_initialized:
            pygame.mixer.stop()

        # Clean up playback thread safely
        if self.playback_thread and self.playback_thread.is_alive():
            self.playback_thread.join(timeout=1.0)

        # Clear the buffers
        with self.buffer_lock:
            self.audio_buffer.clear()
            self.mark_queue.clear()
            self.buffered_samples = 0
            self.samples_played = 0

    def get_state(self) -> Dict[str, Any]:
        """Get current audio output state for debugging"""
        return {
            "is_playing": self.is_playing,
            "is_buffering": self.is_buffering,
            "pygame_initialized": self.pygame_initialized,
            "playback_thread_alive": self.playback_thread.is_alive() if self.playback_thread else False,
            "audio_buffer_size": len(self.audio_buffer),
            "buffered_seconds": self._get_buffered_seconds(),
            "pygame_busy": pygame.mixer.get_busy() if self.pygame_initialized else False,
            "echo_buffer_size": len(self.echo_buffer)
        }

    def get_echo_buffer(self) -> List[np.ndarray]:
        """Get current echo buffer for cancellation"""
        return self.echo_buffer.copy()

    @staticmethod
    def _resample_audio(
        audio_data: np.ndarray,
        original_rate: int,
        target_rate: int
    ) -> np.ndarray:
        """Resample audio data to target sample rate."""
        if original_rate == target_rate:
            return audio_data

        # Calculate resampling factor
        resample_factor = target_rate / original_rate
        original_length = len(audio_data)
        target_length = int(original_length * resample_factor)

        # Use numpy interpolation for better quality
        x_original = np.arange(original_length)
        x_target = np.linspace(0, original_length - 1, target_length)
        resampled = np.interp(x_target, x_original, audio_data)

        return resampled.astype(np.float32)
