using System.Runtime.CompilerServices;
using System.Threading.Channels;

namespace MeetingAssistant.Recording;

public sealed class CompositeMeetingAudioSource : IMeetingAudioSource
{
    private readonly IMeetingAudioSource microphone;
    private readonly IMeetingAudioSource systemAudio;
    private readonly ILogger<CompositeMeetingAudioSource> logger;

    public CompositeMeetingAudioSource(
        IMeetingAudioSource microphone,
        IMeetingAudioSource systemAudio,
        ILogger<CompositeMeetingAudioSource> logger)
    {
        this.microphone = microphone;
        this.systemAudio = systemAudio;
        this.logger = logger;
    }

    public async IAsyncEnumerable<AudioChunk> CaptureAsync(
        [EnumeratorCancellation] CancellationToken cancellationToken)
    {
        var channel = Channel.CreateUnbounded<SourceChunk>();
        var microphoneTask = PumpAsync(AudioSourceKind.Microphone, microphone, channel.Writer, cancellationToken);
        var systemTask = PumpAsync(AudioSourceKind.System, systemAudio, channel.Writer, cancellationToken);
        _ = Task.WhenAll(microphoneTask, systemTask)
            .ContinueWith(
                task =>
                {
                    if (task.Exception is not null)
                    {
                        channel.Writer.TryComplete(task.Exception);
                        return;
                    }

                    channel.Writer.TryComplete();
                },
                CancellationToken.None,
                TaskContinuationOptions.ExecuteSynchronously,
                TaskScheduler.Default);

        AudioChunk? pendingMicrophone = null;
        AudioChunk? pendingSystem = null;

        await foreach (var sourceChunk in channel.Reader.ReadAllAsync(cancellationToken))
        {
            logger.LogDebug(
                "Captured {ByteCount} byte(s) from {SourceKind}",
                sourceChunk.Chunk.Pcm.Length,
                sourceChunk.Kind);
            if (sourceChunk.Kind == AudioSourceKind.Microphone)
            {
                pendingMicrophone = sourceChunk.Chunk;
            }
            else
            {
                pendingSystem = sourceChunk.Chunk;
            }

            if (pendingMicrophone is not null && pendingSystem is not null)
            {
                yield return Mix(pendingMicrophone, pendingSystem);
                pendingMicrophone = null;
                pendingSystem = null;
                continue;
            }

            using var timeout = new CancellationTokenSource(TimeSpan.FromMilliseconds(50));
            using var linked = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, timeout.Token);
            try
            {
                if (await channel.Reader.WaitToReadAsync(linked.Token))
                {
                    continue;
                }
            }
            catch (OperationCanceledException) when (timeout.IsCancellationRequested && !cancellationToken.IsCancellationRequested)
            {
            }

            if (pendingMicrophone is not null)
            {
                yield return pendingMicrophone;
                pendingMicrophone = null;
            }

            if (pendingSystem is not null)
            {
                yield return pendingSystem;
                pendingSystem = null;
            }
        }

        if (pendingMicrophone is not null)
        {
            yield return pendingMicrophone;
        }

        if (pendingSystem is not null)
        {
            yield return pendingSystem;
        }
    }

    private static async Task PumpAsync(
        AudioSourceKind kind,
        IMeetingAudioSource source,
        ChannelWriter<SourceChunk> writer,
        CancellationToken cancellationToken)
    {
        try
        {
            await foreach (var chunk in source.CaptureAsync(cancellationToken))
            {
                await writer.WriteAsync(new SourceChunk(kind, chunk), cancellationToken);
            }
        }
        catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
        {
            throw;
        }
        catch (Exception exception)
        {
            writer.TryComplete(new InvalidOperationException($"{kind} audio capture failed.", exception));
            throw;
        }
    }

    private AudioChunk Mix(AudioChunk left, AudioChunk right)
    {
        if (left.SampleRate != right.SampleRate || left.Channels != right.Channels)
        {
            logger.LogWarning(
                "Cannot mix audio chunks with different formats. Using microphone format {MicrophoneSampleRate}/{MicrophoneChannels} and dropping system chunk format {SystemSampleRate}/{SystemChannels}.",
                left.SampleRate,
                left.Channels,
                right.SampleRate,
                right.Channels);
            return left;
        }

        var mixed = new byte[Math.Max(left.Pcm.Length, right.Pcm.Length)];
        var sampleCount = mixed.Length / sizeof(short);

        for (var sampleIndex = 0; sampleIndex < sampleCount; sampleIndex++)
        {
            var byteIndex = sampleIndex * sizeof(short);
            var leftSample = ReadSample(left.Pcm, byteIndex);
            var rightSample = ReadSample(right.Pcm, byteIndex);
            var sample = Math.Clamp(leftSample + rightSample, short.MinValue, short.MaxValue);
            BitConverter.GetBytes((short)sample).CopyTo(mixed, byteIndex);
        }

        return new AudioChunk(mixed, left.SampleRate, left.Channels);
    }

    private static int ReadSample(byte[] pcm, int byteIndex)
    {
        return byteIndex + 1 < pcm.Length ? BitConverter.ToInt16(pcm, byteIndex) : 0;
    }

    private sealed record SourceChunk(AudioSourceKind Kind, AudioChunk Chunk);

    private enum AudioSourceKind
    {
        Microphone,
        System
    }
}