"""Audio recording module with real-time sound level display."""

import queue
import sys
import threading
import time
from collections import deque

import numpy as np
import sounddevice as sd
from rich.console import Console
from rich.live import Live
from rich.progress import BarColumn, Progress, SpinnerColumn, TextColumn


console = Console()


class AudioRecorder:
    """Records audio from microphone with voice activity detection."""
    
    def __init__(self, sample_rate=16000, channels=1, silence_threshold=0.01, silence_duration=3.0):
        """Initialize audio recorder.
        
        Args:
            sample_rate: Audio sample rate in Hz
            channels: Number of audio channels (1 for mono)
            silence_threshold: Threshold for silence detection (RMS value)
            silence_duration: Duration of silence before stopping (seconds)
        """
        self.sample_rate = sample_rate
        self.channels = channels
        self.silence_threshold = silence_threshold
        self.silence_duration = silence_duration
        self.audio_queue = queue.Queue()
        self.is_recording = False
        self.audio_buffer = []
        self.silence_start = None
        
    def _audio_callback(self, indata, frames, time_info, status):
        """Callback for audio stream."""
        if status:
            console.print(f"[red]Audio status: {status}[/red]")
        self.audio_queue.put(indata.copy())

    def _clear_queue(self):
        """Clear the audio queue."""
        while not self.audio_queue.empty():
            try:
                self.audio_queue.get_nowait()
            except queue.Empty:
                break
        
    def _calculate_rms(self, audio_data):
        """Calculate RMS (Root Mean Square) of audio data."""
        return np.sqrt(np.mean(audio_data**2))
        
    def _display_audio_level(self, rms, max_bars=30):
        """Create a visual representation of audio level."""
        # Normalize RMS to 0-1 range (assuming max RMS around 0.3)
        normalized = min(rms / 0.3, 1.0)
        num_bars = int(normalized * max_bars)
        
        # Color based on level
        if normalized < 0.1:
            color = "dim"
        elif normalized < 0.3:
            color = "green"
        elif normalized < 0.6:
            color = "yellow"
        else:
            color = "red"
            
        bar = "█" * num_bars + "░" * (max_bars - num_bars)
        return f"[{color}]{bar}[/{color}] {rms:.4f}"
        
    def record(self, status_callback=None):
        """Record audio until silence timeout.

        Args:
            status_callback: Optional callback function(status) where status is:
                '2' = recording with sound detected
                '1' = recording with silence detected
                Called only when status changes.

        Returns:
            numpy.ndarray: Recorded audio data
        """
        # Clear any leftover audio from previous recordings
        self._clear_queue()

        self.is_recording = True
        self.audio_buffer = []
        self.silence_start = None

        # Buffer for recent audio levels (for display smoothing)
        rms_history = deque(maxlen=10)

        # Track audio state: '2' = sound, '1' = silence, None = initial
        audio_state = None

        console.print("[bold green]Recording started...[/bold green]")
        console.print("[dim]Speak now. Recording will stop after 3 seconds of silence.[/dim]\n")

        with sd.InputStream(
            samplerate=self.sample_rate,
            channels=self.channels,
            dtype=np.float32,
            blocksize=int(self.sample_rate * 0.1),  # 100ms blocks
            callback=self._audio_callback
        ):
            with Live(console=console, refresh_per_second=10) as live:
                while self.is_recording:
                    try:
                        # Get audio data (timeout allows checking is_recording)
                        audio_data = self.audio_queue.get(timeout=0.1)
                        self.audio_buffer.append(audio_data)

                        # Calculate RMS
                        rms = self._calculate_rms(audio_data)
                        rms_history.append(rms)
                        avg_rms = sum(rms_history) / len(rms_history)

                        # Display audio level
                        level_display = self._display_audio_level(avg_rms)
                        live.update(level_display)

                        # Voice activity detection
                        if rms > self.silence_threshold:
                            self.silence_start = None
                            # Sound detected
                            if audio_state != '2':
                                audio_state = '2'
                                if status_callback:
                                    status_callback('2')
                        else:
                            if self.silence_start is None:
                                self.silence_start = time.time()
                                # Silence detected
                                if audio_state != '1':
                                    audio_state = '1'
                                    if status_callback:
                                        status_callback('1')
                            elif time.time() - self.silence_start >= self.silence_duration:
                                console.print("\n[bold yellow]Silence detected - stopping recording[/bold yellow]")
                                self.is_recording = False

                    except queue.Empty:
                        continue

        # Ensure proper cleanup after Live context exits
        sys.stdout.flush()
        time.sleep(0.1)  # Brief pause to let terminal settle

        # Concatenate all audio chunks
        if self.audio_buffer:
            audio_data = np.concatenate(self.audio_buffer, axis=0)
            console.print(f"[green]Recorded {len(audio_data) / self.sample_rate:.1f} seconds of audio[/green]")
            sys.stdout.flush()
            return audio_data
        else:
            console.print("[red]No audio recorded[/red]")
            sys.stdout.flush()
            return None
            
    def stop(self):
        """Stop recording manually."""
        self.is_recording = False

    def calibrate_silence(self, duration=10.0, silence_sample_duration=3.0):
        """Calibrate silence threshold by recording audio.
        
        Records for specified duration and uses the last silence_sample_duration
        seconds to determine an appropriate silence threshold.
        
        Args:
            duration: Total recording duration in seconds (default: 10.0)
            silence_sample_duration: Duration of silence sample at end (default: 3.0)
            
        Returns:
            float: Recommended silence threshold value
        """
        # Clear any leftover audio from previous recordings
        self._clear_queue()

        self.is_recording = True
        self.audio_buffer = []

        console.print(f"[bold cyan]Calibrating silence threshold...[/bold cyan]")
        console.print(f"[dim]Please remain silent for the last {silence_sample_duration:.0f} seconds of the recording.[/dim]")
        console.print(f"[dim]Recording for {duration:.0f} seconds total...[/dim]\n")
        
        start_time = time.time()
        rms_history = deque(maxlen=10)
        
        with sd.InputStream(
            samplerate=self.sample_rate,
            channels=self.channels,
            dtype=np.float32,
            blocksize=int(self.sample_rate * 0.1),  # 100ms blocks
            callback=self._audio_callback
        ):
            with Live(console=console, refresh_per_second=10) as live:
                while self.is_recording and (time.time() - start_time) < duration:
                    try:
                        # Get audio data
                        audio_data = self.audio_queue.get(timeout=0.1)
                        self.audio_buffer.append(audio_data)
                        
                        # Calculate RMS
                        rms = self._calculate_rms(audio_data)
                        rms_history.append(rms)
                        avg_rms = sum(rms_history) / len(rms_history)
                        
                        # Display audio level
                        level_display = self._display_audio_level(avg_rms)
                        
                        # Show countdown
                        remaining = duration - (time.time() - start_time)
                        if remaining <= silence_sample_duration:
                            level_display += f"\n[yellow]Keep silent! ({remaining:.1f}s)[/yellow]"
                        else:
                            level_display += f"\n[dim]Recording... ({remaining:.1f}s)[/dim]"
                        
                        live.update(level_display)
                        
                    except queue.Empty:
                        continue
        
        self.is_recording = False
        
        if not self.audio_buffer:
            console.print("[red]No audio recorded during calibration[/red]")
            return None
        
        # Concatenate all audio
        audio_data = np.concatenate(self.audio_buffer, axis=0)
        
        # Extract the last silence_sample_duration seconds
        silence_samples = int(silence_sample_duration * self.sample_rate)
        if len(audio_data) < silence_samples:
            console.print("[yellow]Warning: Recording too short, using all available audio[/yellow]")
            silence_audio = audio_data
        else:
            silence_audio = audio_data[-silence_samples:]
        
        # Calculate statistics on the silence portion
        # Split into 100ms chunks and calculate RMS for each
        chunk_size = int(0.1 * self.sample_rate)  # 100ms chunks
        rms_values = []
        
        for i in range(0, len(silence_audio), chunk_size):
            chunk = silence_audio[i:i + chunk_size]
            if len(chunk) > 0:
                rms = self._calculate_rms(chunk)
                rms_values.append(rms)
        
        if not rms_values:
            console.print("[red]Could not analyze silence portion[/red]")
            return None
        
        # Calculate statistics
        mean_rms = np.mean(rms_values)
        std_rms = np.std(rms_values)
        max_rms = np.max(rms_values)
        percentile_95 = np.percentile(rms_values, 95)
        
        # Recommended threshold: mean + 3*std, but at least 1.5x the mean
        # Also ensure it's above the 95th percentile with some margin
        recommended_threshold = max(
            mean_rms + 3 * std_rms,
            mean_rms * 1.5,
            percentile_95 * 1.2
        )
        
        # Cap at reasonable values
        recommended_threshold = min(recommended_threshold, 0.1)
        recommended_threshold = max(recommended_threshold, 0.001)
        
        console.print(f"\n[bold green]Calibration complete![/bold green]")
        console.print(f"[dim]Silence statistics (last {silence_sample_duration:.0f}s):[/dim]")
        console.print(f"  Mean RMS: {mean_rms:.6f}")
        console.print(f"  Std RMS:  {std_rms:.6f}")
        console.print(f"  Max RMS:  {max_rms:.6f}")
        console.print(f"  95th percentile: {percentile_95:.6f}")
        console.print(f"\n[bold cyan]Recommended silence threshold: {recommended_threshold:.6f}[/bold cyan]")
        
        return recommended_threshold
