#import "WebRTCFrameCryptorBridge.h"

#import <WebRTC/RTCRtpReceiver.h>
#import <WebRTC/RTCRtpSender.h>

#include <algorithm>
#include <array>
#include <atomic>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <mutex>
#include <string>
#include <unordered_map>
#include <vector>
#include <mach/mach_time.h>
#include <os/log.h>

#include "NativeXChaCha20.hpp"

namespace rtc {

enum class RefCountReleaseStatus {
    kDroppedLastRef,
    kOtherRefsRemained
};

class RefCountInterface {
public:
    virtual void AddRef() const = 0;
    virtual RefCountReleaseStatus Release() const = 0;

protected:
    virtual ~RefCountInterface() = default;
};

template <typename T>
class scoped_refptr {
public:
    scoped_refptr() : ptr_(nullptr) {}

    explicit scoped_refptr(T *ptr) : ptr_(ptr) {
        if (ptr_ != nullptr) {
            ptr_->AddRef();
        }
    }

    scoped_refptr(const scoped_refptr &other) : ptr_(other.ptr_) {
        if (ptr_ != nullptr) {
            ptr_->AddRef();
        }
    }

    scoped_refptr(scoped_refptr &&other) noexcept : ptr_(other.ptr_) {
        other.ptr_ = nullptr;
    }

    ~scoped_refptr() {
        if (ptr_ != nullptr) {
            ptr_->Release();
            ptr_ = nullptr;
        }
    }

    scoped_refptr &operator=(const scoped_refptr &other) {
        if (this == &other) {
            return *this;
        }
        T *next = other.ptr_;
        if (next != nullptr) {
            next->AddRef();
        }
        if (ptr_ != nullptr) {
            ptr_->Release();
        }
        ptr_ = next;
        return *this;
    }

    scoped_refptr &operator=(scoped_refptr &&other) noexcept {
        if (this == &other) {
            return *this;
        }
        if (ptr_ != nullptr) {
            ptr_->Release();
        }
        ptr_ = other.ptr_;
        other.ptr_ = nullptr;
        return *this;
    }

    scoped_refptr &operator=(T *ptr) {
        if (ptr == ptr_) {
            return *this;
        }
        if (ptr != nullptr) {
            ptr->AddRef();
        }
        if (ptr_ != nullptr) {
            ptr_->Release();
        }
        ptr_ = ptr;
        return *this;
    }

    T *get() const { return ptr_; }
    T *operator->() const { return ptr_; }
    explicit operator bool() const { return ptr_ != nullptr; }

private:
    T *ptr_;
};

template <typename T>
inline scoped_refptr<T> make_ref_counted(T *ptr) {
    return scoped_refptr<T>(ptr);
}

template <typename T>
class ArrayView {
public:
    ArrayView() : data_(nullptr), size_(0) {}
    ArrayView(T *data, std::size_t size) : data_(data), size_(size) {}

    T *data() const { return data_; }
    std::size_t size() const { return size_; }

private:
    T *data_;
    std::size_t size_;
};

} // namespace rtc

namespace cricket {

enum MediaType {
    MEDIA_TYPE_AUDIO,
    MEDIA_TYPE_VIDEO,
    MEDIA_TYPE_DATA,
    MEDIA_TYPE_UNSUPPORTED,
    MEDIA_TYPE_ANY
};

} // namespace cricket

namespace webrtc {

class FrameEncryptorInterface : public rtc::RefCountInterface {
public:
    virtual int Encrypt(cricket::MediaType media_type,
                        uint32_t ssrc,
                        rtc::ArrayView<const uint8_t> additional_data,
                        rtc::ArrayView<const uint8_t> frame,
                        rtc::ArrayView<uint8_t> encrypted_frame,
                        std::size_t *bytes_written) = 0;
    virtual std::size_t GetMaxCiphertextByteSize(cricket::MediaType media_type,
                                                 std::size_t frame_size) = 0;

protected:
    ~FrameEncryptorInterface() override = default;
};

class FrameDecryptorInterface : public rtc::RefCountInterface {
public:
    struct Result {
        enum class Status {
            kOk = 0,
            kRecoverable = 1,
            kFailedToDecrypt = 2
        };

        Result(Status s, std::size_t bw) : status(s), bytes_written(bw) {}
        bool IsOk() const { return status == Status::kOk; }

        Status status;
        std::size_t bytes_written;
    };

    virtual Result Decrypt(cricket::MediaType media_type,
                           const std::vector<uint32_t> &csrcs,
                           rtc::ArrayView<const uint8_t> additional_data,
                           rtc::ArrayView<const uint8_t> encrypted_frame,
                           rtc::ArrayView<uint8_t> frame) = 0;
    virtual std::size_t GetMaxPlaintextByteSize(cricket::MediaType media_type,
                                                std::size_t encrypted_frame_size) = 0;

protected:
    ~FrameDecryptorInterface() override = default;
};

} // namespace webrtc

@interface RTCRtpSender (RosettaFrameCryptorPrivate)
- (void)setFrameEncryptor:(rtc::scoped_refptr<webrtc::FrameEncryptorInterface>)encryptor;
@end

@interface RTCRtpReceiver (RosettaFrameCryptorPrivate)
- (void)setFrameDecryptor:(rtc::scoped_refptr<webrtc::FrameDecryptorInterface>)decryptor;
@end

namespace {

constexpr std::size_t kSharedKeyLength = 32;
constexpr std::size_t kNonceLength = 24;
constexpr int kDiagFrameLimit = 50;

static os_log_t e2ee_log() {
    static os_log_t log = os_log_create("com.rosetta.e2ee", "frame");
    return log;
}

std::mutex gCryptorLock;
std::unordered_map<std::string, rtc::scoped_refptr<webrtc::FrameEncryptorInterface>> gSenderEncryptors;
std::unordered_map<std::string, rtc::scoped_refptr<webrtc::FrameDecryptorInterface>> gReceiverDecryptors;

std::string senderMapKey(RTCRtpSender *sender) {
    NSString *senderId = sender.senderId;
    if (senderId.length > 0) {
        return std::string(senderId.UTF8String);
    }
    return "sender@" + std::to_string(reinterpret_cast<std::uintptr_t>(sender));
}

std::string receiverMapKey(RTCRtpReceiver *receiver) {
    NSString *receiverId = receiver.receiverId;
    if (receiverId.length > 0) {
        return std::string(receiverId.UTF8String);
    }
    return "receiver@" + std::to_string(reinterpret_cast<std::uintptr_t>(receiver));
}

// MARK: - RTP helpers (ported from rosetta_e2ee.cpp)

struct ParsedRtpPacket {
    std::size_t header_size = 0;
    uint16_t sequence = 0;
    uint32_t timestamp = 0;
    uint32_t ssrc = 0;
};

struct RtpProbeState {
    bool locked = false;
    bool has_probe = false;
    uint32_t probe_ssrc = 0;
    uint16_t probe_sequence = 0;
    uint32_t probe_timestamp = 0;
    uint32_t ssrc = 0;
    uint16_t last_sequence = 0;
    uint32_t last_timestamp = 0;
};

struct GeneratedTsState {
    bool initialized = false;
    uint32_t next_timestamp = 0;
    uint32_t next_step = 960;
};

struct AdditionalTsState {
    bool initialized = false;
    uint64_t base_timestamp = 0;
};

inline uint16_t load16_be(const uint8_t *p) {
    return (uint16_t)((uint16_t)p[0] << 8 | (uint16_t)p[1]);
}

inline uint32_t load32_be(const uint8_t *p) {
    return (uint32_t)p[0] << 24 | (uint32_t)p[1] << 16 | (uint32_t)p[2] << 8 | (uint32_t)p[3];
}

inline uint64_t load64_be(const uint8_t *p) {
    return (uint64_t)p[0] << 56 | (uint64_t)p[1] << 48 | (uint64_t)p[2] << 40 | (uint64_t)p[3] << 32 |
           (uint64_t)p[4] << 24 | (uint64_t)p[5] << 16 | (uint64_t)p[6] << 8  | (uint64_t)p[7];
}

inline void fill_nonce_from_ts32(uint32_t ts, uint8_t nonce[24]) {
    nonce[4] = (uint8_t)(ts >> 24);
    nonce[5] = (uint8_t)(ts >> 16);
    nonce[6] = (uint8_t)(ts >> 8);
    nonce[7] = (uint8_t)(ts);
}

inline void fill_nonce_from_ts64(uint64_t ts, uint8_t nonce[24]) {
    nonce[0] = (uint8_t)(ts >> 56);
    nonce[1] = (uint8_t)(ts >> 48);
    nonce[2] = (uint8_t)(ts >> 40);
    nonce[3] = (uint8_t)(ts >> 32);
    nonce[4] = (uint8_t)(ts >> 24);
    nonce[5] = (uint8_t)(ts >> 16);
    nonce[6] = (uint8_t)(ts >> 8);
    nonce[7] = (uint8_t)(ts);
}

bool parse_rtp_packet(const uint8_t *data, std::size_t len, ParsedRtpPacket *out) {
    if (!data || !out || len < 12) return false;
    const uint8_t version = (data[0] >> 6) & 0x03;
    if (version != 2) return false;
    const std::size_t csrc_count = (std::size_t)(data[0] & 0x0F);
    const bool has_extension = (data[0] & 0x10) != 0;
    std::size_t header = 12 + csrc_count * 4;
    if (header > len) return false;
    if (has_extension) {
        if (len < header + 4) return false;
        const uint16_t ext_len_words = (uint16_t)((uint16_t)data[header + 2] << 8 | (uint16_t)data[header + 3]);
        header += 4 + (std::size_t)ext_len_words * 4;
        if (header > len) return false;
    }
    const std::size_t payload_size = len - header;
    if (payload_size == 0 || payload_size > 1200) return false;
    out->header_size = header;
    out->sequence = load16_be(data + 2);
    out->timestamp = load32_be(data + 4);
    out->ssrc = load32_be(data + 8);
    return true;
}

bool fill_nonce_from_rtp_frame(const uint8_t *data, std::size_t len,
                               RtpProbeState *state, uint8_t nonce[24],
                               std::size_t *header_size) {
    if (!state) return false;
    ParsedRtpPacket packet;
    if (!parse_rtp_packet(data, len, &packet)) return false;

    if (!state->locked) {
        if (!state->has_probe) {
            state->has_probe = true;
            state->probe_ssrc = packet.ssrc;
            state->probe_sequence = packet.sequence;
            state->probe_timestamp = packet.timestamp;
            return false;
        }
        const bool same_ssrc = packet.ssrc == state->probe_ssrc;
        const uint16_t seq_delta = (uint16_t)(packet.sequence - state->probe_sequence);
        const bool sequence_progressed = seq_delta > 0 && seq_delta <= 10;
        if (!same_ssrc || !sequence_progressed) {
            state->probe_ssrc = packet.ssrc;
            state->probe_sequence = packet.sequence;
            state->probe_timestamp = packet.timestamp;
            return false;
        }
        state->locked = true;
        state->ssrc = packet.ssrc;
        state->last_sequence = packet.sequence;
        state->last_timestamp = packet.timestamp;
    } else {
        if (packet.ssrc != state->ssrc) {
            state->locked = false;
            state->has_probe = true;
            state->probe_ssrc = packet.ssrc;
            state->probe_sequence = packet.sequence;
            state->probe_timestamp = packet.timestamp;
            return false;
        }
        const uint16_t seq_delta = (uint16_t)(packet.sequence - state->last_sequence);
        if (seq_delta != 0 && seq_delta <= 200) {
            state->last_sequence = packet.sequence;
            state->last_timestamp = packet.timestamp;
        } else if (seq_delta != 0) {
            state->locked = false;
            state->has_probe = true;
            state->probe_ssrc = packet.ssrc;
            state->probe_sequence = packet.sequence;
            state->probe_timestamp = packet.timestamp;
            return false;
        }
    }
    nonce[4] = (uint8_t)(packet.timestamp >> 24);
    nonce[5] = (uint8_t)(packet.timestamp >> 16);
    nonce[6] = (uint8_t)(packet.timestamp >> 8);
    nonce[7] = (uint8_t)(packet.timestamp);
    if (header_size) *header_size = packet.header_size;
    return true;
}

bool fill_nonce_from_additional_data(const uint8_t *data, std::size_t len,
                                     uint8_t nonce[24],
                                     bool *used_rtp_header,
                                     AdditionalTsState *ts_state,
                                     bool *used_relative_ts) {
    if (used_rtp_header) *used_rtp_header = false;
    if (used_relative_ts) *used_relative_ts = false;
    if (!data || len < 8) return false;

    if (len == 8) {
        uint64_t ts64 = load64_be(data);
        uint64_t ts_rel = ts64;
        if (ts_state != nullptr) {
            if (!ts_state->initialized) {
                ts_state->initialized = true;
                ts_state->base_timestamp = ts64;
            } else if (ts64 < ts_state->base_timestamp) {
                ts_state->base_timestamp = ts64;
            }
            ts_rel = ts64 - ts_state->base_timestamp;
            if (used_relative_ts) *used_relative_ts = true;
        }
        fill_nonce_from_ts64(ts_rel, nonce);
        return true;
    }

    if (len >= 12) {
        const uint8_t version = (data[0] >> 6) & 0x03;
        if (version == 2) {
            uint32_t ts = load32_be(data + 4);
            fill_nonce_from_ts32(ts, nonce);
            if (used_rtp_header) *used_rtp_header = true;
            return true;
        }
    }

    std::memcpy(nonce, data, std::min<std::size_t>(8, len));
    return true;
}

// MARK: - Opus plausibility checks (ported from rosetta_e2ee.cpp)

uint32_t opus_base_frame_samples(uint8_t config) {
    if (config <= 11) {
        static const uint32_t kSilk[4] = {480, 960, 1920, 2880};
        return kSilk[config & 0x03];
    }
    if (config <= 15) {
        return (config & 0x01) ? 960 : 480;
    }
    static const uint32_t kCelt[4] = {120, 240, 480, 960};
    return kCelt[config & 0x03];
}

uint32_t infer_opus_packet_duration_samples(const uint8_t *packet, std::size_t len) {
    if (!packet || len == 0) return 960;
    const uint8_t toc = packet[0];
    const uint8_t config = (uint8_t)(toc >> 3);
    const uint8_t frame_code = (uint8_t)(toc & 0x03);
    uint32_t frame_count = 1;
    if (frame_code == 1 || frame_code == 2) {
        frame_count = 2;
    } else if (frame_code == 3) {
        if (len < 2) return 960;
        frame_count = (uint32_t)(packet[1] & 0x3F);
        if (frame_count == 0 || frame_count > 48) return 960;
    }
    uint32_t base = opus_base_frame_samples(config);
    uint32_t total = base * frame_count;
    if (total < 120 || total > 5760) return 960;
    return total;
}

bool decode_opus_frame_len(const uint8_t *packet, std::size_t len,
                           std::size_t *offset, std::size_t *frame_len) {
    if (!packet || !offset || !frame_len) return false;
    if (*offset >= len) return false;
    const uint8_t b0 = packet[*offset];
    (*offset)++;
    if (b0 < 252) { *frame_len = b0; return true; }
    if (*offset >= len) return false;
    const uint8_t b1 = packet[*offset];
    (*offset)++;
    *frame_len = (std::size_t)b1 * 4u + (std::size_t)b0;
    return true;
}

bool is_plausible_opus_packet(const uint8_t *packet, std::size_t len) {
    if (!packet || len == 0 || len > 2000) return false;
    const uint8_t toc = packet[0];
    const uint8_t config = (uint8_t)(toc >> 3);
    if (config > 31) return false;
    const uint8_t frame_code = (uint8_t)(toc & 0x03);
    const std::size_t payload_len = len - 1;
    if (frame_code == 0) return payload_len >= 1;
    if (frame_code == 1) { if (payload_len < 2) return false; return (payload_len % 2) == 0; }
    if (frame_code == 2) {
        std::size_t off = 1;
        std::size_t len1 = 0;
        if (!decode_opus_frame_len(packet, len, &off, &len1)) return false;
        if (len1 == 0) return false;
        if (off + len1 >= len) return false;
        return (len - off - len1) > 0;
    }
    if (len < 2) return false;
    const uint8_t ch = packet[1];
    const bool has_padding = (ch & 0x40) != 0;
    const uint8_t frame_count = (uint8_t)(ch & 0x3F);
    if (frame_count == 0 || frame_count > 48) return false;
    uint32_t total = opus_base_frame_samples(config) * (uint32_t)frame_count;
    if (total > 5760) return false;
    if (has_padding) return false;
    const bool vbr = (ch & 0x80) != 0;
    if (!vbr) {
        const std::size_t data_len = len - 2;
        if (data_len < (std::size_t)frame_count) return false;
        return (data_len % frame_count) == 0;
    }
    std::size_t off = 2;
    std::size_t consumed = 0;
    for (std::size_t i = 0; i + 1 < frame_count; ++i) {
        std::size_t flen = 0;
        if (!decode_opus_frame_len(packet, len, &off, &flen)) return false;
        if (flen == 0) return false;
        consumed += flen;
        if (off + consumed >= len) return false;
    }
    return (len - off) > consumed;
}

bool is_plausible_decrypted_audio_frame(const uint8_t *data, std::size_t len) {
    if (!data || len == 0) return false;
    ParsedRtpPacket packet;
    if (parse_rtp_packet(data, len, &packet) && packet.header_size < len) {
        return is_plausible_opus_packet(data + packet.header_size, len - packet.header_size);
    }
    return is_plausible_opus_packet(data, len);
}

// MARK: - Encryptor (Android-parity with multi-strategy nonce)

class RosettaFrameEncryptor final : public webrtc::FrameEncryptorInterface {
public:
    explicit RosettaFrameEncryptor(const std::array<std::uint8_t, kSharedKeyLength> &key) {
        std::memcpy(key_.data(), key.data(), kSharedKeyLength);
    }

    int Encrypt(cricket::MediaType media_type,
                uint32_t ssrc,
                rtc::ArrayView<const uint8_t> additional_data,
                rtc::ArrayView<const uint8_t> frame,
                rtc::ArrayView<uint8_t> encrypted_frame,
                std::size_t *bytes_written) override {
        (void)media_type;
        (void)ssrc;

        if (frame.size() == 0 || encrypted_frame.size() < frame.size()) {
            *bytes_written = 0;
            return -1;
        }

        uint8_t nonce[24] = {0};
        bool ad_was_rtp = false;
        const char *mode = "zero";

        // Strategy 1: additional_data (Desktop-compatible 8-byte timestamp).
        // After custom WebRTC build with E2EE patch, ad=8 on every frame.
        bool has_nonce = fill_nonce_from_additional_data(
            additional_data.data(), additional_data.size(),
            nonce, &ad_was_rtp, nullptr, nullptr);
        if (has_nonce) {
            mode = ad_was_rtp ? "ad-rtp" : "raw-abs";
        }

        // Strategy 2: RTP header in frame payload
        if (!has_nonce) {
            std::size_t header_size = 0;
            has_nonce = fill_nonce_from_rtp_frame(
                frame.data(), frame.size(), &rtp_probe_, nonce, &header_size);
            if (has_nonce) {
                mode = "rtp";
            }
        }

        // Strategy 3: nonce stays zero (fallback for unpatched WebRTC builds).
        // Both sender and receiver will use nonce=0 symmetrically.

        // Always encrypt — E2EE is mandatory for all call frames.
        std::vector<std::uint8_t> encrypted;
        const bool ok = rosetta::nativecrypto::xchacha20_xor(
            frame.data(), frame.size(), key_.data(), nonce, encrypted);
        if (!ok || encrypted.size() != frame.size()) {
            *bytes_written = 0;
            return 1;
        }
        std::memcpy(encrypted_frame.data(), encrypted.data(), frame.size());
        *bytes_written = frame.size();

        // Diagnostic log for first N frames
        int n = diag_count_.fetch_add(1, std::memory_order_relaxed);
        if (n < kDiagFrameLimit) {
            os_log_info(e2ee_log(),
                "ENC frame#%d sz=%zu ad=%zu mode=%s nonce=[%02x%02x%02x%02x%02x%02x%02x%02x]",
                n, frame.size(), additional_data.size(), mode,
                nonce[0], nonce[1], nonce[2], nonce[3],
                nonce[4], nonce[5], nonce[6], nonce[7]);
        }

        return 0;
    }

    std::size_t GetMaxCiphertextByteSize(cricket::MediaType media_type, std::size_t frame_size) override {
        (void)media_type;
        return frame_size;
    }

    void AddRef() const override {
        refCount_.fetch_add(1, std::memory_order_relaxed);
    }

    rtc::RefCountReleaseStatus Release() const override {
        const int previous = refCount_.fetch_sub(1, std::memory_order_acq_rel);
        if (previous == 1) {
            delete this;
            return rtc::RefCountReleaseStatus::kDroppedLastRef;
        }
        return rtc::RefCountReleaseStatus::kOtherRefsRemained;
    }

private:
    mutable std::atomic<int> refCount_{0};
    mutable std::atomic<int> diag_count_{0};
    std::array<std::uint8_t, kSharedKeyLength> key_;
    mutable RtpProbeState rtp_probe_;
};

// MARK: - Decryptor (with plausibility checks when timestamp available)

class RosettaFrameDecryptor final : public webrtc::FrameDecryptorInterface {
public:
    explicit RosettaFrameDecryptor(const std::array<std::uint8_t, kSharedKeyLength> &key) {
        std::memcpy(key_.data(), key.data(), kSharedKeyLength);
    }

    Result Decrypt(cricket::MediaType media_type,
                   const std::vector<uint32_t> &csrcs,
                   rtc::ArrayView<const uint8_t> additional_data,
                   rtc::ArrayView<const uint8_t> encrypted_frame,
                   rtc::ArrayView<uint8_t> frame) override {
        (void)media_type;
        (void)csrcs;

        if (encrypted_frame.size() == 0 || frame.size() < encrypted_frame.size()) {
            return Result(Result::Status::kFailedToDecrypt, 0);
        }

        uint8_t nonce[24] = {0};
        bool ad_was_rtp = false;
        const char *mode = "zero";
        bool plausible = true;

        // Strategy 1: additional_data (Desktop-compatible 8-byte timestamp).
        // After custom WebRTC build with E2EE patch, ad=8 on every frame.
        bool has_nonce = fill_nonce_from_additional_data(
            additional_data.data(), additional_data.size(),
            nonce, &ad_was_rtp, nullptr, nullptr);
        if (has_nonce) {
            mode = ad_was_rtp ? "ad-rtp" : "raw-abs";
        }

        // Strategy 2: RTP header in encrypted frame
        if (!has_nonce) {
            std::size_t header_size = 0;
            has_nonce = fill_nonce_from_rtp_frame(
                encrypted_frame.data(), encrypted_frame.size(),
                &rtp_probe_, nonce, &header_size);
            if (has_nonce) {
                mode = "rtp";
            }
        }

        // Strategy 3: nonce stays zero (fallback for unpatched WebRTC builds).

        // Always decrypt — E2EE is mandatory for all call frames.
        std::vector<std::uint8_t> decrypted;
        rosetta::nativecrypto::xchacha20_xor(
            encrypted_frame.data(), encrypted_frame.size(),
            key_.data(), nonce, decrypted);
        if (decrypted.size() != encrypted_frame.size()) {
            return Result(Result::Status::kFailedToDecrypt, 0);
        }
        std::memcpy(frame.data(), decrypted.data(), encrypted_frame.size());

        plausible = is_plausible_decrypted_audio_frame(frame.data(), encrypted_frame.size());

        // Fallback: try relative timestamp (when ad=8 but base differs)
        if (!plausible && has_nonce && additional_data.size() == 8 && !ad_was_rtp) {
            uint8_t nonce_rel[24] = {0};
            bool rel_used = false;
            if (fill_nonce_from_additional_data(
                    additional_data.data(), additional_data.size(),
                    nonce_rel, nullptr, &additional_rel_ts_state_, &rel_used) && rel_used) {
                std::vector<std::uint8_t> candidate;
                rosetta::nativecrypto::xchacha20_xor(
                    encrypted_frame.data(), encrypted_frame.size(),
                    key_.data(), nonce_rel, candidate);
                if (is_plausible_decrypted_audio_frame(candidate.data(), candidate.size())) {
                    std::memcpy(frame.data(), candidate.data(), candidate.size());
                    std::memcpy(nonce, nonce_rel, sizeof(nonce));
                    plausible = true;
                    mode = "raw-rel-fb";
                }
            }
        }

        // Passthrough: peer sent unencrypted frames (raw Opus)
        if (!plausible && has_nonce &&
            is_plausible_decrypted_audio_frame(encrypted_frame.data(), encrypted_frame.size())) {
            std::memcpy(frame.data(), encrypted_frame.data(), encrypted_frame.size());
            plausible = true;
            mode = "passthrough";
        }

        int n = diag_count_.fetch_add(1, std::memory_order_relaxed);
        if (n < kDiagFrameLimit) {
            os_log_info(e2ee_log(),
                "DEC frame#%d sz=%zu ad=%zu mode=%s ok=%d nonce=[%02x%02x%02x%02x%02x%02x%02x%02x]",
                n, encrypted_frame.size(), additional_data.size(), mode,
                plausible ? 1 : 0,
                nonce[0], nonce[1], nonce[2], nonce[3],
                nonce[4], nonce[5], nonce[6], nonce[7]);
        }

        return Result(Result::Status::kOk, encrypted_frame.size());
    }

    std::size_t GetMaxPlaintextByteSize(cricket::MediaType media_type,
                                        std::size_t encrypted_frame_size) override {
        (void)media_type;
        return encrypted_frame_size;
    }

    void AddRef() const override {
        refCount_.fetch_add(1, std::memory_order_relaxed);
    }

    rtc::RefCountReleaseStatus Release() const override {
        const int previous = refCount_.fetch_sub(1, std::memory_order_acq_rel);
        if (previous == 1) {
            delete this;
            return rtc::RefCountReleaseStatus::kDroppedLastRef;
        }
        return rtc::RefCountReleaseStatus::kOtherRefsRemained;
    }

private:
    mutable std::atomic<int> refCount_{0};
    mutable std::atomic<int> diag_count_{0};
    std::array<std::uint8_t, kSharedKeyLength> key_;
    mutable RtpProbeState rtp_probe_;
    mutable AdditionalTsState additional_rel_ts_state_;
};

// MARK: - Key normalization

std::array<std::uint8_t, kSharedKeyLength> normalizeSharedKey(NSData *sharedKey) {
    std::array<std::uint8_t, kSharedKeyLength> key{};
    if (sharedKey.length == 0 || sharedKey.bytes == nullptr) {
        return key;
    }
    const std::size_t copyLength = std::min<std::size_t>(kSharedKeyLength, sharedKey.length);
    std::memcpy(key.data(), sharedKey.bytes, copyLength);
    return key;
}

} // namespace

@implementation WebRTCFrameCryptorBridge

+ (BOOL)attachSender:(RTCRtpSender *)sender sharedKey:(NSData *)sharedKey {
    if (sender == nil || sharedKey == nil || sharedKey.length < kSharedKeyLength) {
        return NO;
    }
    if (![sender respondsToSelector:@selector(setFrameEncryptor:)]) {
        return NO;
    }

    const std::string key = senderMapKey(sender);
    std::lock_guard<std::mutex> lock(gCryptorLock);
    const auto existing = gSenderEncryptors.find(key);
    if (existing != gSenderEncryptors.end()) {
        [sender setFrameEncryptor:existing->second];
        return YES;
    }

    const auto normalizedKey = normalizeSharedKey(sharedKey);
    rtc::scoped_refptr<webrtc::FrameEncryptorInterface> encryptor(
        new RosettaFrameEncryptor(normalizedKey)
    );

    [sender setFrameEncryptor:encryptor];
    gSenderEncryptors[key] = encryptor;
    return YES;
}

+ (BOOL)attachReceiver:(RTCRtpReceiver *)receiver sharedKey:(NSData *)sharedKey {
    if (receiver == nil || sharedKey == nil || sharedKey.length < kSharedKeyLength) {
        return NO;
    }
    if (![receiver respondsToSelector:@selector(setFrameDecryptor:)]) {
        return NO;
    }

    const std::string key = receiverMapKey(receiver);
    std::lock_guard<std::mutex> lock(gCryptorLock);
    const auto existing = gReceiverDecryptors.find(key);
    if (existing != gReceiverDecryptors.end()) {
        [receiver setFrameDecryptor:existing->second];
        return YES;
    }

    const auto normalizedKey = normalizeSharedKey(sharedKey);
    rtc::scoped_refptr<webrtc::FrameDecryptorInterface> decryptor(
        new RosettaFrameDecryptor(normalizedKey)
    );

    [receiver setFrameDecryptor:decryptor];
    gReceiverDecryptors[key] = decryptor;
    return YES;
}

+ (void)detachSender:(RTCRtpSender *)sender {
    if (sender == nil) {
        return;
    }

    std::lock_guard<std::mutex> lock(gCryptorLock);
    const std::string key = senderMapKey(sender);
    if ([sender respondsToSelector:@selector(setFrameEncryptor:)]) {
        [sender setFrameEncryptor:rtc::scoped_refptr<webrtc::FrameEncryptorInterface>()];
    }
    gSenderEncryptors.erase(key);
}

+ (void)detachReceiver:(RTCRtpReceiver *)receiver {
    if (receiver == nil) {
        return;
    }

    std::lock_guard<std::mutex> lock(gCryptorLock);
    const std::string key = receiverMapKey(receiver);
    if ([receiver respondsToSelector:@selector(setFrameDecryptor:)]) {
        [receiver setFrameDecryptor:rtc::scoped_refptr<webrtc::FrameDecryptorInterface>()];
    }
    gReceiverDecryptors.erase(key);
}

@end
