using MeetingAssistant.Recording;
using Microsoft.Extensions.Options;
using NAudio.Wave;
using Whisper.net;

namespace MeetingAssistant.Transcription;

public sealed class WhisperLocalStreamingTranscriptionProvider : IStreamingTranscriptionProvider
{
    private readonly MeetingAssistantOptions options;
    private readonly ILogger<WhisperLocalStreamingTranscriptionProvider> logger;

    public WhisperLocalStreamingTranscriptionProvider(
        IOptions<MeetingAssistantOptions> options,
        ILogger<WhisperLocalStreamingTranscriptionProvider> logger)
    {
        this.options = options.Value;
        this.logger = logger;
    }

    public async IAsyncEnumerable<TranscriptionSegment> TranscribeAsync(
        IAsyncEnumerable<AudioChunk> audio,
        SpeechRecognitionPipelineOptions pipelineOptions,
        [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken cancellationToken)
    {
        var modelPath = Path.GetFullPath(Environment.ExpandEnvironmentVariables(options.WhisperLocal.ModelPath));
        if (!File.Exists(modelPath))
        {
            throw new FileNotFoundException($"Whisper model was not found at '{modelPath}'.", modelPath);
        }

        using var factory = WhisperFactory.FromPath(modelPath);
        using var processor = CreateProcessor(factory, pipelineOptions.DictationWords);

        var window = new List<AudioChunk>();
        var windowDuration = TimeSpan.Zero;
        var windowTarget = TimeSpan.FromSeconds(Math.Max(1, options.WhisperLocal.WindowSeconds));
        var offset = TimeSpan.Zero;

        logger.LogInformation(
            "Starting local Whisper streaming transcription with model {ModelPath} and {WindowSeconds}s windows",
            modelPath,
            windowTarget.TotalSeconds);

        await foreach (var chunk in audio.WithCancellation(cancellationToken))
        {
            window.Add(chunk);
            windowDuration += chunk.Duration;

            if (windowDuration >= windowTarget)
            {
                var segments = 0;
                await foreach (var segment in ProcessWindow(processor, window, offset, cancellationToken))
                {
                    segments++;
                    logger.LogInformation(
                        "Whisper emitted segment at {Start}: {Text}",
                        segment.Start,
                        segment.Text);
                    yield return segment;
                }

                logger.LogInformation(
                    "Processed Whisper window at offset {Offset} with duration {Duration} and {SegmentCount} segment(s)",
                    offset,
                    windowDuration,
                    segments);
                offset += windowDuration;
                window.Clear();
                windowDuration = TimeSpan.Zero;
            }
        }

        if (window.Count > 0)
        {
            var segments = 0;
            await foreach (var segment in ProcessWindow(processor, window, offset, cancellationToken))
            {
                segments++;
                logger.LogInformation(
                    "Whisper emitted segment at {Start}: {Text}",
                    segment.Start,
                    segment.Text);
                yield return segment;
            }

            logger.LogInformation(
                "Processed final Whisper window at offset {Offset} with duration {Duration} and {SegmentCount} segment(s)",
                offset,
                windowDuration,
                segments);
        }
    }

    private async IAsyncEnumerable<TranscriptionSegment> ProcessWindow(
        WhisperProcessor processor,
        IReadOnlyList<AudioChunk> chunks,
        TimeSpan offset,
        [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken cancellationToken)
    {
        var stream = CreateWaveStream(chunks);

        await foreach (var result in processor.ProcessAsync(stream, cancellationToken))
        {
            var text = result.Text?.Trim();
            if (string.IsNullOrWhiteSpace(text))
            {
                continue;
            }

            if (IsControlMarker(text))
            {
                logger.LogInformation("Whisper emitted control marker {Marker}; skipping transcript output", text);
                continue;
            }

            yield return new TranscriptionSegment(offset + result.Start, offset + result.End, "Unknown", text);
        }
    }

    private WhisperProcessor CreateProcessor(
        WhisperFactory factory,
        IReadOnlyList<string>? dictationWords)
    {
        var builder = factory.CreateBuilder()
            .WithLanguage(options.WhisperLocal.Language);
        var prompt = BuildPrompt(dictationWords);
        if (!string.IsNullOrWhiteSpace(prompt))
        {
            builder.WithPrompt(prompt)
                .WithCarryInitialPrompt(true);
            logger.LogInformation(
                "Configured Whisper prompt with {WordCount} dictation word(s)",
                dictationWords?.Count ?? 0);
        }

        return builder.Build();
    }

    internal static string BuildPrompt(IReadOnlyList<string>? dictationWords)
    {
        var words = (dictationWords ?? [])
            .Select(word => word.Trim())
            .Where(word => !string.IsNullOrWhiteSpace(word))
            .Distinct(StringComparer.OrdinalIgnoreCase)
            .ToList();
        return words.Count == 0
            ? ""
            : "Prefer these domain terms and acronyms: " + string.Join(", ", words) + ".";
    }

    private static bool IsControlMarker(string text)
    {
        return text is "[BLANK_AUDIO]" or "[MUSIC]" or "[NO_SPEECH]" or "[SILENCE]";
    }

    private MemoryStream CreateWaveStream(IReadOnlyList<AudioChunk> chunks)
    {
        var first = chunks[0];
        var stream = new MemoryStream();
        var pcm = new MemoryStream();

        foreach (var chunk in chunks)
        {
            if (chunk.SampleRate != first.SampleRate || chunk.Channels != first.Channels)
            {
                logger.LogWarning("Dropping audio chunk with mismatched format while creating Whisper window");
                continue;
            }

            pcm.Write(chunk.Pcm, 0, chunk.Pcm.Length);
        }

        WriteWaveHeader(stream, first.SampleRate, first.Channels, (int)pcm.Length);
        pcm.Position = 0;
        pcm.CopyTo(stream);
        stream.Position = 0;
        return stream;
    }

    private static void WriteWaveHeader(Stream stream, int sampleRate, int channels, int dataLength)
    {
        using var writer = new BinaryWriter(stream, System.Text.Encoding.ASCII, leaveOpen: true);
        writer.Write(System.Text.Encoding.ASCII.GetBytes("RIFF"));
        writer.Write(36 + dataLength);
        writer.Write(System.Text.Encoding.ASCII.GetBytes("WAVE"));
        writer.Write(System.Text.Encoding.ASCII.GetBytes("fmt "));
        writer.Write(16);
        writer.Write((short)1);
        writer.Write((short)channels);
        writer.Write(sampleRate);
        writer.Write(sampleRate * channels * sizeof(short));
        writer.Write((short)(channels * sizeof(short)));
        writer.Write((short)16);
        writer.Write(System.Text.Encoding.ASCII.GetBytes("data"));
        writer.Write(dataLength);
    }
}