SoundFlix/FractalVisualizer/Audio/AudioCapture.cs

using System;
using System.Threading.Tasks;
using FractalVisualizer.Models;
using NAudio.CoreAudioApi;
using NAudio.CoreAudioApi.Interfaces;
using NAudio.Wave;

namespace FractalVisualizer.Audio;

public sealed class AudioCapture : IDisposable
{
    private readonly object _stateSync = new();

    private readonly MMDeviceEnumerator _deviceEnumerator;
    private readonly EndpointNotificationClient _notificationClient;

    private IWaveIn? _capture;
    private bool _isRunning;
    private bool _isRestarting;
    private bool _disposed;

    public AudioCapture(int ringBufferCapacity = 48000 * 8)
    {
        RingBuffer = new AudioSampleRingBuffer(ringBufferCapacity);
        CurrentMode = AudioInputMode.Microphone;
        SampleRate = 48000;

        _deviceEnumerator = new MMDeviceEnumerator();
        _notificationClient = new EndpointNotificationClient(this);
    }

    public AudioSampleRingBuffer RingBuffer { get; }

    public AudioInputMode CurrentMode { get; private set; }

    public int SampleRate { get; private set; }

    public event Action<AudioInputMode>? ModeChanged;

    public event Action<int>? SampleRateChanged;

    public void Start()
    {
        ThrowIfDisposed();

        lock (_stateSync)
        {
            if (_isRunning)
            {
                return;
            }

            _deviceEnumerator.RegisterEndpointNotificationCallback(_notificationClient);
            _isRunning = true;
            StartCaptureForCurrentModeNoLock(clearBuffer: true);
        }
    }

    public void Stop()
    {
        lock (_stateSync)
        {
            if (!_isRunning)
            {
                return;
            }

            _isRunning = false;
            StopAndDisposeCaptureNoLock();
            _deviceEnumerator.UnregisterEndpointNotificationCallback(_notificationClient);
        }
    }

    public void SwitchMode(AudioInputMode newMode)
    {
        ThrowIfDisposed();

        lock (_stateSync)
        {
            if (CurrentMode == newMode)
            {
                return;
            }

            CurrentMode = newMode;
            if (_isRunning)
            {
                RestartCaptureNoLock(clearBuffer: true);
            }
            else
            {
                RingBuffer.Clear();
            }
        }

        ModeChanged?.Invoke(CurrentMode);
    }

    private void HandleDefaultDeviceChanged(DataFlow flow, Role role)
    {
        if (role != Role.Multimedia)
        {
            return;
        }

        lock (_stateSync)
        {
            if (!_isRunning || _isRestarting)
            {
                return;
            }

            if (CurrentMode == AudioInputMode.Microphone && flow == DataFlow.Capture)
            {
                RestartCaptureNoLock(clearBuffer: false);
                return;
            }

            if (CurrentMode == AudioInputMode.SystemOutput && flow == DataFlow.Render)
            {
                RestartCaptureNoLock(clearBuffer: false);
            }
        }
    }

    private void RestartCaptureNoLock(bool clearBuffer)
    {
        _isRestarting = true;
        try
        {
            StopAndDisposeCaptureNoLock();
            StartCaptureForCurrentModeNoLock(clearBuffer);
        }
        finally
        {
            _isRestarting = false;
        }
    }

    private void StartCaptureForCurrentModeNoLock(bool clearBuffer)
    {
        if (clearBuffer)
        {
            RingBuffer.Clear();
        }

        _capture = CurrentMode switch
        {
            AudioInputMode.Microphone => CreateMicrophoneCapture(),
            AudioInputMode.SystemOutput => CreateLoopbackCapture(),
            _ => throw new ArgumentOutOfRangeException(nameof(CurrentMode), CurrentMode, "Unknown audio input mode.")
        };

        _capture.DataAvailable += OnDataAvailable;
        _capture.RecordingStopped += OnRecordingStopped;
        _capture.StartRecording();

        SampleRate = _capture.WaveFormat.SampleRate;
        SampleRateChanged?.Invoke(SampleRate);
    }

    private IWaveIn CreateMicrophoneCapture()
    {
        try
        {
            MMDevice defaultInput = _deviceEnumerator.GetDefaultAudioEndpoint(DataFlow.Capture, Role.Multimedia);
            return new WasapiCapture(defaultInput);
        }
        catch
        {
            return new WaveInEvent
            {
                WaveFormat = new WaveFormat(44100, 16, 1),
                BufferMilliseconds = 25
            };
        }
    }

    private IWaveIn CreateLoopbackCapture()
    {
        try
        {
            MMDevice defaultOutput = _deviceEnumerator.GetDefaultAudioEndpoint(DataFlow.Render, Role.Multimedia);
            return new WasapiLoopbackCapture(defaultOutput);
        }
        catch
        {
            return new WasapiLoopbackCapture();
        }
    }

    private void OnDataAvailable(object? sender, WaveInEventArgs e)
    {
        IWaveIn? activeCapture = _capture;
        if (activeCapture is null || e.BytesRecorded <= 0)
        {
            return;
        }

        float[] samples = ConvertToMonoFloatSamples(e.Buffer, e.BytesRecorded, activeCapture.WaveFormat);
        if (samples.Length == 0)
        {
            return;
        }

        RingBuffer.Write(samples);
    }

    private static float[] ConvertToMonoFloatSamples(byte[] buffer, int bytesRecorded, WaveFormat format)
    {
        int channels = Math.Max(1, format.Channels);
        if (format.Encoding == WaveFormatEncoding.IeeeFloat && format.BitsPerSample == 32)
        {
            int totalSamples = bytesRecorded / sizeof(float);
            int frames = totalSamples / channels;
            if (frames <= 0)
            {
                return Array.Empty<float>();
            }

            float[] mono = new float[frames];
            WaveBuffer waveBuffer = new(buffer);
            for (int frame = 0; frame < frames; frame++)
            {
                float sum = 0f;
                int offset = frame * channels;
                for (int ch = 0; ch < channels; ch++)
                {
                    sum += waveBuffer.FloatBuffer[offset + ch];
                }

                mono[frame] = sum / channels;
            }

            return mono;
        }

        if (format.Encoding == WaveFormatEncoding.Pcm && format.BitsPerSample == 16)
        {
            int totalSamples = bytesRecorded / sizeof(short);
            int frames = totalSamples / channels;
            if (frames <= 0)
            {
                return Array.Empty<float>();
            }

            float[] mono = new float[frames];
            WaveBuffer waveBuffer = new(buffer);
            for (int frame = 0; frame < frames; frame++)
            {
                float sum = 0f;
                int offset = frame * channels;
                for (int ch = 0; ch < channels; ch++)
                {
                    sum += waveBuffer.ShortBuffer[offset + ch] / 32768f;
                }

                mono[frame] = sum / channels;
            }

            return mono;
        }

        return Array.Empty<float>();
    }

    private void OnRecordingStopped(object? sender, StoppedEventArgs e)
    {
        if (_disposed)
        {
            return;
        }

        if (e.Exception is null)
        {
            return;
        }

        Task.Run(() =>
        {
            lock (_stateSync)
            {
                if (!_isRunning || _isRestarting)
                {
                    return;
                }

                RestartCaptureNoLock(clearBuffer: false);
            }
        });
    }

    private void StopAndDisposeCaptureNoLock()
    {
        if (_capture is null)
        {
            return;
        }

        IWaveIn toDispose = _capture;
        _capture = null;

        toDispose.DataAvailable -= OnDataAvailable;
        toDispose.RecordingStopped -= OnRecordingStopped;

        try
        {
            toDispose.StopRecording();
        }
        catch
        {
            // Device may already be removed/disposed.
        }

        toDispose.Dispose();
    }

    private void ThrowIfDisposed()
    {
        if (_disposed)
        {
            throw new ObjectDisposedException(nameof(AudioCapture));
        }
    }

    public void Dispose()
    {
        if (_disposed)
        {
            return;
        }

        Stop();
        _deviceEnumerator.Dispose();
        _disposed = true;
    }

    private sealed class EndpointNotificationClient(AudioCapture owner) : IMMNotificationClient
    {
        private readonly AudioCapture _owner = owner;

        public void OnDefaultDeviceChanged(DataFlow flow, Role role, string defaultDeviceId)
        {
            _owner.HandleDefaultDeviceChanged(flow, role);
        }

        public void OnDeviceAdded(string pwstrDeviceId)
        {
        }

        public void OnDeviceRemoved(string deviceId)
        {
        }

        public void OnDeviceStateChanged(string deviceId, DeviceState newState)
        {
        }

        public void OnPropertyValueChanged(string pwstrDeviceId, PropertyKey key)
        {
        }
    }
}