using VoiceCat.Protocol;

namespace VoiceCat.Crypto;

public sealed class MediaDecryptor : IDisposable
{
    private readonly MediaCipher cipher;
    private ulong highestSequence;
    private ulong replayWindow;
    private bool initialized;
    private bool disposed;

    public MediaDecryptor(ReadOnlySpan<byte> key) : this(key, false) { }

    internal MediaDecryptor(ReadOnlySpan<byte> key, bool useManaged) => cipher = new(key, useManaged);

    public bool TryDecrypt(ReadOnlySpan<byte> packet, Span<byte> plaintext, out VoiceFrameHeader header, out int bytesWritten)
    {
        ObjectDisposedException.ThrowIf(disposed, this);
        header = default;
        bytesWritten = 0;
        if (packet.Length < VoiceFrameHeader.Size + MediaEncryptor.TagSize) return false;
        int length = packet.Length - VoiceFrameHeader.Size - MediaEncryptor.TagSize;
        ArgumentOutOfRangeException.ThrowIfLessThan(plaintext.Length, length);
        if (packet.Overlaps(plaintext)) throw new ArgumentException("Input and output must not overlap.", nameof(plaintext));
        VoiceFrameHeader.TryRead(packet, out var candidate);
        ulong sequence = candidate.Sequence;
        if (initialized && sequence <= highestSequence)
        {
            ulong offset = highestSequence - sequence;
            if (offset >= 64 || (replayWindow & (1UL << (int)offset)) != 0) return false;
        }
        if (!cipher.TryDecrypt(sequence, packet[VoiceFrameHeader.Size..], packet[..VoiceFrameHeader.Size], plaintext[..length])) return false;

        // Only authenticated counters may move the replay window.
        if (!initialized)
        {
            highestSequence = sequence;
            replayWindow = 1;
            initialized = true;
        }
        else if (sequence > highestSequence)
        {
            ulong shift = sequence - highestSequence;
            replayWindow = (shift >= 64 ? 0 : replayWindow << (int)shift) | 1;
            highestSequence = sequence;
        }
        else replayWindow |= 1UL << (int)(highestSequence - sequence);
        header = candidate;
        bytesWritten = length;
        return true;
    }

    public void Dispose()
    {
        if (disposed) return;
        disposed = true;
        cipher.Dispose();
    }
}