"""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."""
        normalized = min(rms / 0.3, 1.0)
        num_bars = int(normalized * max_bars)

        if rms < self.silence_threshold:
            color = "cyan"
        elif 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
        """
        self._clear_queue()

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

        rms_history = deque(maxlen=10)

        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),
            callback=self._audio_callback
        ):
            with Live(console=console, refresh_per_second=10) as live:
                while self.is_recording:
                    try:
                        audio_data = self.audio_queue.get(timeout=0.1)
                        self.audio_buffer.append(audio_data)

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

                        level_display = self._display_audio_level(avg_rms)
                        live.update(level_display)

                        if rms > self.silence_threshold:
                            self.silence_start = None
                            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()
                                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

        sys.stdout.flush()
        time.sleep(0.1)

        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.
        
        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
        """
        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),
            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:
                        audio_data = self.audio_queue.get(timeout=0.1)
                        self.audio_buffer.append(audio_data)
                        
                        rms = self._calculate_rms(audio_data)
                        rms_history.append(rms)
                        avg_rms = sum(rms_history) / len(rms_history)
                        
                        level_display = self._display_audio_level(avg_rms)
                        
                        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
        
        audio_data = np.concatenate(self.audio_buffer, axis=0)
        
        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:]
        
        chunk_size = int(0.1 * self.sample_rate)
        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
        
        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 = max(
            mean_rms + 3 * std_rms,
            mean_rms * 1.5,
            percentile_95 * 1.2
        )
        
        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

    def measure_noise_quick(self, duration=5.0, skip_first=1.0):
        """Quickly measure ambient noise level.
        
        Args:
            duration: Total recording duration in seconds (default: 5.0)
            skip_first: Seconds to skip at the beginning (default: 1.0)
            
        Returns:
            float: Recommended silence threshold value based on measured noise
        """
        self._clear_queue()
        self.is_recording = True
        self.audio_buffer = []

        console.print(f"[bold cyan]Measuring ambient noise level...[/bold cyan]")
        console.print(f"[dim]Please remain silent for {duration:.0f} seconds (skipping first {skip_first:.0f}s)...[/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),
            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:
                        audio_data = self.audio_queue.get(timeout=0.1)
                        self.audio_buffer.append(audio_data)
                        
                        rms = self._calculate_rms(audio_data)
                        rms_history.append(rms)
                        avg_rms = sum(rms_history) / len(rms_history)
                        
                        level_display = self._display_audio_level(avg_rms)
                        remaining = duration - (time.time() - start_time)
                        elapsed = time.time() - start_time
                        if elapsed < skip_first:
                            level_display += f"\n[yellow]Skipping initial spike... ({skip_first - elapsed:.1f}s)[/yellow]"
                        else:
                            level_display += f"\n[dim]Measuring noise... ({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 measurement[/red]")
            return None
        
        audio_data = np.concatenate(self.audio_buffer, axis=0)
        
        # Skip the first 'skip_first' seconds to avoid initial spike
        skip_samples = int(skip_first * self.sample_rate)
        if len(audio_data) <= skip_samples:
            console.print("[red]Recording too short[/red]")
            return None
        
        audio_data = audio_data[skip_samples:]
        
        # Calculate RMS for remaining audio
        chunk_size = int(0.1 * self.sample_rate)
        rms_values = []
        
        for i in range(0, len(audio_data), chunk_size):
            chunk = audio_data[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 audio[/red]")
            return None
        
        mean_rms = np.mean(rms_values)
        std_rms = np.std(rms_values)
        percentile_95 = np.percentile(rms_values, 95)
        
        recommended_threshold = max(
            mean_rms + 3 * std_rms,
            mean_rms * 1.5,
            percentile_95 * 1.2
        )
        
        recommended_threshold = min(recommended_threshold, 0.1)
        recommended_threshold = max(recommended_threshold, 0.0001)
        
        console.print(f"\n[bold green]Noise measurement complete![/bold green]")
        console.print(f"[dim]Noise statistics (after skipping first {skip_first:.0f}s):[/dim]")
        console.print(f"  Mean RMS: {mean_rms:.6f}")
        console.print(f"  Std RMS:  {std_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
