using MeetingAssistant.Transcription;
using NAudio.Wave;

namespace MeetingAssistant.Speakers;

public interface ISpeakerIdentityDiarizationClient
{
    Task<IReadOnlyList<TranscriptionSegment>> DiarizeAsync(
        string wavPath,
        CancellationToken cancellationToken);
}

public sealed class AzureSpeechSpeakerIdentityMatcher : ISpeakerIdentityMatcher
{
    private readonly ISpeakerIdentityDiarizationClient diarizationClient;
    private readonly SpeakerIdentificationOptions options;
    private readonly ILogger<AzureSpeechSpeakerIdentityMatcher> logger;

    public AzureSpeechSpeakerIdentityMatcher(
        ISpeakerIdentityDiarizationClient diarizationClient,
        Microsoft.Extensions.Options.IOptions<MeetingAssistantOptions> options,
        ILogger<AzureSpeechSpeakerIdentityMatcher> logger)
    {
        this.diarizationClient = diarizationClient;
        this.options = options.Value.SpeakerIdentification;
        this.logger = logger;
    }

    public async Task<SpeakerIdentityMatch?> MatchAsync(
        SpeakerIdentityMatchRequest request,
        CancellationToken cancellationToken)
    {
        if (!options.Enabled || request.Candidates.Count == 0)
        {
            return null;
        }

        logger.LogInformation(
            "Speaker identity matching {DiarizedSpeaker} against {CandidateCount} candidate(s)",
            request.DiarizedSpeaker,
            request.Candidates.Count);
        var tempPath = Path.Combine(Path.GetTempPath(), "meeting-assistant-speaker-match", $"{Guid.NewGuid():N}.wav");
        try
        {
            Directory.CreateDirectory(Path.GetDirectoryName(tempPath)!);
            var layout = WriteCompositeWav(tempPath, request);
            if (layout.UnknownSegment is null || layout.KnownSegments.Count == 0)
            {
                logger.LogInformation(
                    "Speaker identity matching {DiarizedSpeaker} skipped because no readable snippets were available",
                    request.DiarizedSpeaker);
                return null;
            }

            var segments = await DiarizeWithTimeoutAsync(tempPath, cancellationToken);
            if (segments.Count == 0)
            {
                logger.LogInformation(
                    "Speaker identity matching {DiarizedSpeaker} skipped because diarization returned no segments",
                    request.DiarizedSpeaker);
                return null;
            }

            logger.LogInformation(
                "Speaker identity matching {DiarizedSpeaker} received {SegmentCount} diarized segment(s)",
                request.DiarizedSpeaker,
                segments.Count);
            var unknownSpeaker = FindBestSpeaker(layout.UnknownSegment.Value, segments);
            if (string.IsNullOrWhiteSpace(unknownSpeaker))
            {
                logger.LogInformation(
                    "Speaker identity matching {DiarizedSpeaker} found no speaker for the unknown snippet",
                    request.DiarizedSpeaker);
                return null;
            }

            foreach (var known in layout.KnownSegments.GroupBy(segment => segment.IdentityId))
            {
                var matchingSnippetCount = known.Count(segment =>
                    string.Equals(
                        FindBestSpeaker(segment.Segment, segments),
                        unknownSpeaker,
                        StringComparison.Ordinal));
                var requiredMatches = known.Count() > 1 ? 2 : 1;
                if (matchingSnippetCount >= requiredMatches)
                {
                    logger.LogInformation(
                        "Speaker identity matching {DiarizedSpeaker} matched identity {IdentityId}",
                        request.DiarizedSpeaker,
                        known.Key);
                    return new SpeakerIdentityMatch(known.Key);
                }
            }

            logger.LogInformation(
                "Speaker identity matching {DiarizedSpeaker} found no matching identity",
                request.DiarizedSpeaker);
            return null;
        }
        finally
        {
            TryDelete(tempPath);
        }
    }

    private async Task<IReadOnlyList<TranscriptionSegment>> DiarizeWithTimeoutAsync(
        string wavPath,
        CancellationToken cancellationToken)
    {
        var matchTimeout = options.MatchTimeout;
        if (matchTimeout <= TimeSpan.Zero)
        {
            return await diarizationClient.DiarizeAsync(wavPath, cancellationToken);
        }

        using var timeout = new CancellationTokenSource(matchTimeout);
        using var linked = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, timeout.Token);
        try
        {
            return await diarizationClient.DiarizeAsync(wavPath, linked.Token);
        }
        catch (OperationCanceledException) when (timeout.IsCancellationRequested && !cancellationToken.IsCancellationRequested)
        {
            logger.LogWarning(
                "Speaker identity diarization timed out after {Timeout}",
                matchTimeout);
            return [];
        }
    }

    private CompositeLayout WriteCompositeWav(string path, SpeakerIdentityMatchRequest request)
    {
        using var firstReader = OpenFirstReadableWave(request);
        if (firstReader is null)
        {
            return new CompositeLayout(null, []);
        }

        var silence = TimeSpan.FromSeconds(Math.Max(0, options.SilenceBetweenSnippetsSeconds));
        using var writer = new WaveFileWriter(path, firstReader.WaveFormat);
        var current = TimeSpan.Zero;
        var knownSegments = new List<KnownCompositeSegment>();

        foreach (var candidate in request.Candidates)
        {
            foreach (var snippet in candidate.Snippets.Where(storedSnippet => storedSnippet.Length > 0))
            {
                var segment = AppendSnippet(writer, firstReader.WaveFormat, snippet, current);
                if (segment is null)
                {
                    continue;
                }

                knownSegments.Add(new KnownCompositeSegment(candidate.IdentityId, segment.Value));
                current = segment.Value.End + silence;
                WriteSilence(writer, firstReader.WaveFormat, silence);
            }
        }

        var unknownSegment = AppendSnippet(writer, firstReader.WaveFormat, request.UnknownSnippet, current);
        return new CompositeLayout(unknownSegment, knownSegments);
    }

    private static WaveFileReader? OpenFirstReadableWave(SpeakerIdentityMatchRequest request)
    {
        foreach (var bytes in request.Candidates.SelectMany(candidate => candidate.Snippets).Append(request.UnknownSnippet))
        {
            if (bytes.Length == 0)
            {
                continue;
            }

            try
            {
                return new WaveFileReader(new MemoryStream(bytes));
            }
            catch (FormatException)
            {
            }
        }

        return null;
    }

    private static TimeSegment? AppendSnippet(
        WaveFileWriter writer,
        WaveFormat targetFormat,
        byte[] snippet,
        TimeSpan start)
    {
        if (snippet.Length == 0)
        {
            return null;
        }

        using var reader = new WaveFileReader(new MemoryStream(snippet));
        if (!WaveFormatsMatch(reader.WaveFormat, targetFormat))
        {
            throw new InvalidDataException("Speaker identity snippets must use the same WAV format.");
        }

        reader.CopyTo(writer);
        return new TimeSegment(start, start + reader.TotalTime);
    }

    private static void WriteSilence(WaveFileWriter writer, WaveFormat format, TimeSpan duration)
    {
        if (duration <= TimeSpan.Zero)
        {
            return;
        }

        var bytes = new byte[(int)(format.AverageBytesPerSecond * duration.TotalSeconds)];
        writer.Write(bytes, 0, bytes.Length);
    }

    private static string? FindBestSpeaker(TimeSegment segment, IReadOnlyList<TranscriptionSegment> segments)
    {
        var bestOverlap = 0d;
        string? bestSpeaker = null;
        foreach (var transcriptSegment in segments)
        {
            var overlap = Math.Min(segment.End.TotalSeconds, transcriptSegment.End.TotalSeconds)
                - Math.Max(segment.Start.TotalSeconds, transcriptSegment.Start.TotalSeconds);
            if (overlap > bestOverlap)
            {
                bestOverlap = overlap;
                bestSpeaker = transcriptSegment.Speaker;
            }
        }

        return bestSpeaker;
    }

    private static bool WaveFormatsMatch(WaveFormat left, WaveFormat right)
    {
        return left.Encoding == right.Encoding
            && left.SampleRate == right.SampleRate
            && left.Channels == right.Channels
            && left.BitsPerSample == right.BitsPerSample;
    }

    private static void TryDelete(string path)
    {
        try
        {
            if (File.Exists(path))
            {
                File.Delete(path);
            }
        }
        catch (IOException)
        {
        }
    }

    private sealed record CompositeLayout(TimeSegment? UnknownSegment, IReadOnlyList<KnownCompositeSegment> KnownSegments);

    private sealed record KnownCompositeSegment(int IdentityId, TimeSegment Segment);

    private readonly record struct TimeSegment(TimeSpan Start, TimeSpan End);
}