/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package org.mozilla.gecko.media;

import android.annotation.SuppressLint;
import android.media.MediaCodec;
import android.media.MediaCodec.BufferInfo;
import android.media.MediaCodec.CryptoInfo;
import android.os.Build;
import android.os.Parcel;
import android.os.Parcelable;
import androidx.annotation.ChecksSdkIntAtLeast;
import java.lang.reflect.Field;
import java.nio.ByteBuffer;
import org.mozilla.gecko.annotation.WrapForJNI;

// Parcelable carrying input/output sample data and info cross process.
public final class Sample implements Parcelable {
  public static final Sample EOS;

  static {
    final BufferInfo eosInfo = new BufferInfo();
    EOS = new Sample();
    EOS.info.set(0, 0, Long.MIN_VALUE, MediaCodec.BUFFER_FLAG_END_OF_STREAM);
  }

  @WrapForJNI public long session;

  public static final int NO_BUFFER = -1;

  public int bufferId = NO_BUFFER;
  @WrapForJNI public BufferInfo info = new BufferInfo();
  public CryptoInfo cryptoInfo;

  // Simple Linked list for recycling objects.
  // Used to nodify Sample objects. Do not marshal/unmarshal.
  private Sample mNext;
  private static Sample sPool = new Sample();
  private static int sPoolSize = 1;

  private Sample() {}

  private void readInfo(final Parcel in) {
    final int offset = in.readInt();
    final int size = in.readInt();
    final long pts = in.readLong();
    final int flags = in.readInt();

    info.set(offset, size, pts, flags);
  }

  private void readCrypto(final Parcel in) {
    final int hasCryptoInfo = in.readInt();
    if (hasCryptoInfo == 0) {
      cryptoInfo = null;
      return;
    }

    final byte[] iv = in.createByteArray();
    final byte[] key = in.createByteArray();
    final int mode = in.readInt();
    final int[] numBytesOfClearData = in.createIntArray();
    final int[] numBytesOfEncryptedData = in.createIntArray();
    final int numSubSamples = in.readInt();

    if (cryptoInfo == null) {
      cryptoInfo = new CryptoInfo();
    }
    cryptoInfo.set(numSubSamples, numBytesOfClearData, numBytesOfEncryptedData, key, iv, mode);
    if (supportsCryptoPattern()) {
      final int numEncryptBlocks = in.readInt();
      final int numSkipBlocks = in.readInt();
      cryptoInfo.setPattern(new CryptoInfo.Pattern(numEncryptBlocks, numSkipBlocks));
    }
  }

  public Sample set(final BufferInfo info, final CryptoInfo cryptoInfo) {
    setBufferInfo(info);
    setCryptoInfo(cryptoInfo);
    return this;
  }

  public void setBufferInfo(final BufferInfo info) {
    this.info.set(0, info.size, info.presentationTimeUs, info.flags);
  }

  public void setCryptoInfo(final CryptoInfo crypto) {
    if (crypto == null) {
      cryptoInfo = null;
      return;
    }

    if (cryptoInfo == null) {
      cryptoInfo = new CryptoInfo();
    }
    cryptoInfo.set(
        crypto.numSubSamples,
        crypto.numBytesOfClearData,
        crypto.numBytesOfEncryptedData,
        crypto.key,
        crypto.iv,
        crypto.mode);
    if (supportsCryptoPattern()) {
      final CryptoInfo.Pattern pattern = getCryptoPatternCompat(crypto);
      if (pattern == null) {
        return;
      }
      cryptoInfo.setPattern(pattern);
    }
  }

  @WrapForJNI
  public void dispose() {
    if (isEOS()) {
      return;
    }

    bufferId = NO_BUFFER;
    info.set(0, 0, 0, 0);
    if (cryptoInfo != null) {
      cryptoInfo.set(0, null, null, null, null, 0);
    }

    // Recycle it.
    synchronized (CREATOR) {
      this.mNext = sPool;
      sPool = this;
      sPoolSize++;
    }
  }

  public boolean isEOS() {
    return (this == EOS) || ((info.flags & MediaCodec.BUFFER_FLAG_END_OF_STREAM) != 0);
  }

  public static Sample obtain() {
    synchronized (CREATOR) {
      Sample s = null;
      if (sPoolSize > 0) {
        s = sPool;
        sPool = s.mNext;
        s.mNext = null;
        sPoolSize--;
      } else {
        s = new Sample();
      }
      return s;
    }
  }

  public static final Creator<Sample> CREATOR =
      new Creator<Sample>() {
        @Override
        public Sample createFromParcel(final Parcel in) {
          return obtainSample(in);
        }

        @Override
        public Sample[] newArray(final int size) {
          return new Sample[size];
        }

        private Sample obtainSample(final Parcel in) {
          final Sample s = obtain();
          s.session = in.readLong();
          s.bufferId = in.readInt();
          s.readInfo(in);
          s.readCrypto(in);
          return s;
        }
      };

  @Override
  public int describeContents() {
    return 0;
  }

  @Override
  public void writeToParcel(final Parcel dest, final int parcelableFlags) {
    dest.writeLong(session);
    dest.writeInt(bufferId);
    writeInfo(dest);
    writeCrypto(dest);
  }

  private void writeInfo(final Parcel dest) {
    dest.writeInt(info.offset);
    dest.writeInt(info.size);
    dest.writeLong(info.presentationTimeUs);
    dest.writeInt(info.flags);
  }

  private void writeCrypto(final Parcel dest) {
    if (cryptoInfo != null) {
      dest.writeInt(1);
      dest.writeByteArray(cryptoInfo.iv);
      dest.writeByteArray(cryptoInfo.key);
      dest.writeInt(cryptoInfo.mode);
      dest.writeIntArray(cryptoInfo.numBytesOfClearData);
      dest.writeIntArray(cryptoInfo.numBytesOfEncryptedData);
      dest.writeInt(cryptoInfo.numSubSamples);
      if (supportsCryptoPattern()) {
        final CryptoInfo.Pattern pattern = getCryptoPatternCompat(cryptoInfo);
        if (pattern != null) {
          dest.writeInt(pattern.getEncryptBlocks());
          dest.writeInt(pattern.getSkipBlocks());
        } else {
          // Couldn't get pattern - write default values
          dest.writeInt(0);
          dest.writeInt(0);
        }
      }
    } else {
      dest.writeInt(0);
    }
  }

  public static byte[] byteArrayFromBuffer(
      final ByteBuffer buffer, final int offset, final int size) {
    if (buffer == null || buffer.capacity() == 0 || size == 0) {
      return null;
    }
    if (buffer.hasArray() && offset == 0 && buffer.array().length == size) {
      return buffer.array();
    }
    final int length = Math.min(offset + size, buffer.capacity()) - offset;
    final byte[] bytes = new byte[length];
    buffer.position(offset);
    buffer.get(bytes);
    return bytes;
  }

  @Override
  public String toString() {
    if (isEOS()) {
      return "EOS sample";
    }

    final StringBuilder str = new StringBuilder();
    str.append("{ session#:")
        .append(session)
        .append(", buffer#")
        .append(bufferId)
        .append(", info=")
        .append("{ offset=")
        .append(info.offset)
        .append(", size=")
        .append(info.size)
        .append(", pts=")
        .append(info.presentationTimeUs)
        .append(", flags=")
        .append(Integer.toHexString(info.flags))
        .append(" }")
        .append(" }");
    return str.toString();
  }

  @ChecksSdkIntAtLeast(api = android.os.Build.VERSION_CODES.N)
  public static boolean supportsCryptoPattern() {
    return Build.VERSION.SDK_INT >= 24;
  }

  @SuppressLint("DiscouragedPrivateApi")
  public static CryptoInfo.Pattern getCryptoPatternCompat(final CryptoInfo cryptoInfo) {
    if (!supportsCryptoPattern()) {
      return null;
    }
    // getPattern() added in API 31:
    // https://developer.android.com/reference/android/media/MediaCodec.CryptoInfo#getPattern()
    if (Build.VERSION.SDK_INT >= 31) {
      return cryptoInfo.getPattern();
    }

    // CryptoInfo.Pattern added in API 24:
    // https://developer.android.com/reference/android/media/MediaCodec.CryptoInfo.Pattern
    if (Build.VERSION.SDK_INT >= 24) {
      try {
        // Without getPattern(), no way to access the pattern without reflection.
        // https://cs.android.com/android/platform/superproject/+/android-11.0.0_r1:frameworks/base/media/java/android/media/MediaCodec.java;l=2718;drc=3c715d5778e15dc84082e63dc65b382d31fe8e45
        final Field patternField = CryptoInfo.class.getDeclaredField("pattern");
        patternField.setAccessible(true);
        return (CryptoInfo.Pattern) patternField.get(cryptoInfo);
      } catch (final NoSuchFieldException | IllegalAccessException e) {
        return null;
      }
    }
    return null;
  }
}