/*
 * Copyright (c) 2016-present The ZLMediaKit project authors. All Rights Reserved.
 *
 * This file is part of ZLMediaKit(https://github.com/ZLMediaKit/ZLMediaKit).
 *
 * Use of this source code is governed by MIT-like license that can be found in the
 * LICENSE file in the root of the source tree. All contributing project authors
 * may be found in the AUTHORS file in the root of the source tree.
 */

#include "VP9Rtp.h"
#include "Extension/Frame.h"
#include "Common/config.h"

namespace mediakit{

const int16_t kNoPictureId = -1;
const int8_t  kNoTl0PicIdx = -1;
const uint8_t kNoTemporalIdx = 0xFF;
const int kNoKeyIdx = -1;

struct VP9ResolutionLayer {
    int width;
    int height;
};

struct RTPPayloadVP9 {
    bool hasPictureID = false;
    bool interPicturePrediction = false;
    bool hasLayerIndices = false;
    bool flexibleMode = false;
    bool beginningOfLayerFrame = false;
    bool endingOfLayerFrame = false;
    bool hasScalabilityStructure = false;
    bool largePictureID = false;
    int pictureID = -1;
    int temporalID = -1;
    bool isSwitchingUp = false;
    int spatialID = -1;
    bool isInterLayeredDepUsed = false;
    int tl0PicIdx = -1;
    int referenceIdx = -1;
    bool additionalReferenceIdx = false;
    int spatialLayers = -1;
    bool hasResolution = false;
    bool hasGof = false;
    int numberOfFramesInGof = -1;
    std::vector<VP9ResolutionLayer> resolutions;

    int parse(const unsigned char *data, int dataLength);
    bool keyFrame() const { return beginningOfLayerFrame && !interPicturePrediction; }
    std::string dump() const {
        char line[64] = {0};
        snprintf(line, sizeof(line), "%c%c%c%c%c%c%c- %d %d, %d %d",
            hasPictureID ? 'I' : ' ', 
            interPicturePrediction ? 'P' : ' ', 
            hasLayerIndices ? 'L' : ' ', 
            flexibleMode ? 'F' : ' ',
            beginningOfLayerFrame ? 'B' : ' ', 
            endingOfLayerFrame ? 'E' : ' ', 
            hasScalabilityStructure ? 'V' : ' ',
            pictureID, tl0PicIdx,
            spatialID, temporalID);
        return line;
    }
};
//
// VP9 format:
//
// Payload descriptor (Flexible mode F = 1)
//        0 1 2 3 4 5 6 7
//       +-+-+-+-+-+-+-+-+
//       |I|P|L|F|B|E|V|-| (REQUIRED)
//       +-+-+-+-+-+-+-+-+
//  I:   |M| PICTURE ID  | (REQUIRED)
//       +-+-+-+-+-+-+-+-+
//  M:   | EXTENDED PID  | (RECOMMENDED)
//       +-+-+-+-+-+-+-+-+
//  L:   |  T  |U|  S  |D| (CONDITIONALLY RECOMMENDED)
//       +-+-+-+-+-+-+-+-+                             -
//  P,F: | P_DIFF      |N| (CONDITIONALLY REQUIRED)    - up to 3 times
//       +-+-+-+-+-+-+-+-+                             -
//  V:   | SS            |
//       | ..            |
//       +-+-+-+-+-+-+-+-+
//
// Payload descriptor (Non flexible mode F = 0)
//
//       0 1 2 3 4 5 6 7
//      +-+-+-+-+-+-+-+-+
//      |I|P|L|F|B|E|V|-| (REQUIRED)
//      +-+-+-+-+-+-+-+-+
// I:   |M| PICTURE ID  | (RECOMMENDED)
//      +-+-+-+-+-+-+-+-+
// M:   | EXTENDED PID  | (RECOMMENDED)
//      +-+-+-+-+-+-+-+-+
// L:   |  T  |U|  S  |D| (CONDITIONALLY RECOMMENDED)
//      +-+-+-+-+-+-+-+-+
//      |   TL0PICIDX   | (CONDITIONALLY REQUIRED)
//      +-+-+-+-+-+-+-+-+
// V:   | SS            |
//      | ..            |
//      +-+-+-+-+-+-+-+-+
#define kIBit 0x80
#define kPBit 0x40
#define kLBit 0x20
#define kFBit 0x10
#define kBBit 0x08
#define kEBit 0x04
#define kVBit 0x02
int RTPPayloadVP9::parse(const unsigned char *data, int dataLength) {
  if (!data || dataLength <= 0) {
    return -1;
  }
  const unsigned char *dataPtr = data;
  size_t remaining = static_cast<size_t>(dataLength);
  auto readByte = [&]() -> int {
    if (remaining < 1) {
      return -1;
    }
    --remaining;
    return *dataPtr++;
  };
  auto skipBytes = [&](size_t count) -> bool {
    if (remaining < count) {
      return false;
    }
    remaining -= count;
    dataPtr += count;
    return true;
  };

  // Parse mandatory first byte of payload descriptor
  int byte = readByte();
  if (byte < 0) return -1;
  this->hasPictureID = (byte & kIBit); // I bit
  this->interPicturePrediction = (byte & kPBit); // P bit
  this->hasLayerIndices = (byte & kLBit); // L bit
  this->flexibleMode = (byte & kFBit); // F bit
  this->beginningOfLayerFrame = (byte & kBBit); // B bit
  this->endingOfLayerFrame = (byte & kEBit); // E bit
  this->hasScalabilityStructure = (byte & kVBit); // V bit

  if (this->hasPictureID) {
    byte = readByte();
    if (byte < 0) return -1;
    this->largePictureID = (byte & 0x80);  // M bit
    this->pictureID = (byte & 0x7F);
    if (this->largePictureID) {
      byte = readByte();
      if (byte < 0) return -1;
      this->pictureID = (this->pictureID << 8) | byte;
    }
  }

  if (this->hasLayerIndices) {
    byte = readByte();
    if (byte < 0) return -1;
    this->temporalID = (byte & 0xE0) >> 5;  // T bits
    this->isSwitchingUp = (byte & 0x10);  // U bit
    this->spatialID = (byte & 0x0E) >> 1;  // S bits
    this->isInterLayeredDepUsed = (byte & 0x01);  // D bit
    if (!this->flexibleMode) {
      byte = readByte();
      if (byte < 0) return -1;
      this->tl0PicIdx = byte;
    }
  }

  if (this->flexibleMode && this->interPicturePrediction) {
      // The VP9 payload descriptor permits at most three reference indices.
      bool nbit;
      int referenceCount = 0;
      do {
          if (++referenceCount > 3) return -1;
          byte = readByte();
          if (byte < 0) return -1;
          this->referenceIdx = (byte & 0xFE) >> 1;
          nbit = (byte & 0x01);
          this->additionalReferenceIdx = nbit;
      } while (nbit);
  }
  if (this->hasScalabilityStructure) {
    byte = readByte();
    if (byte < 0) return -1;
    this->spatialLayers = (byte & 0xE0) >> 5;  // N_S bits
    this->hasResolution = (byte & 0x10);  // Y bit
    this->hasGof = (byte & 0x08);  // G bit
    if (this->hasResolution) {
      for (int i = 0; i <= this->spatialLayers; i++) {
        if (remaining < 4) return -1;
        int width = (dataPtr[0] << 8) + dataPtr[1];
        int height = (dataPtr[2] << 8) + dataPtr[3];
        this->resolutions.push_back({ width, height });
        skipBytes(4);
      }
    }
    if (this->hasGof) {
      byte = readByte();
      if (byte < 0) return -1;
      this->numberOfFramesInGof = byte;  // N_G bits
      for (int frame_index = 0; frame_index < this->numberOfFramesInGof; frame_index++) {
        // TODO(javierc): Read these values if needed
        byte = readByte();
        if (byte < 0) return -1;
        int reference_indices = (byte & 0x0C) >> 2;  // R bits
        if (!skipBytes(reference_indices)) return -1;
      }
    }
  }
  if (remaining == 0) return -1;
  return dataPtr - data;
}


////////////////////////////////////////////////////
VP9RtpDecoder::VP9RtpDecoder() {
    obtainFrame();
}

void VP9RtpDecoder::obtainFrame() {
    _frame = FrameImp::create<VP9Frame>();
}

bool VP9RtpDecoder::inputRtp(const RtpPacket::Ptr &rtp, bool key_pos) {
    auto seq = rtp->getSeq();
    bool is_gop = decodeRtp(rtp);
    if (!_gop_dropped && seq != (uint16_t)(_last_seq + 1) && _last_seq) {
        _gop_dropped = true;
        WarnL << "start drop VP9 gop, last seq:" << _last_seq << ", rtp:\r\n" << rtp->dumpString();
    }
    _last_seq = seq;
    return is_gop;
}

bool VP9RtpDecoder::decodeRtp(const RtpPacket::Ptr &rtp) {
    auto payload_size = rtp->getPayloadSize();
    if (payload_size < 1) {
        // No actual payload
        return false;
    }
    auto payload = rtp->getPayload();
    auto stamp = rtp->getStampMS();
    auto seq = rtp->getSeq();

    RTPPayloadVP9 info;
    int offset = info.parse(payload, payload_size);
    if (offset < 0 || static_cast<size_t>(offset) >= payload_size) {
        WarnL << "VP9 RTP payload parse failed, seq:" << seq;
        _frame_drop = true;
        _frame->_buffer.clear();
        return false;
    }
    // InfoL << rtp->dumpString() << "\n" << info.dump();
    bool start = info.beginningOfLayerFrame;
    if (start) {
        _frame->_pts = stamp;
        _frame->_buffer.clear();
        _frame_drop = false;
    }

    if (_frame_drop) {
        // This frame is incomplete
        return false;
    }

    if (!start && seq != (uint16_t)(_last_seq + 1)) {
        // 中间的或末尾的rtp包，其seq必须连续，否则说明rtp丢包，那么该帧不完整，必须得丢弃
        _frame_drop = true;
        _frame->_buffer.clear();
        return false;
    }
    // Append data
    _frame->_buffer.append((char *)payload + offset, payload_size - offset);
    if (info.endingOfLayerFrame) { // rtp->getHeader()->mark
        // 确保下一个包必须是beginningOfLayerFrame
        _frame_drop = true;
        // 该帧最后一个rtp包,输出frame
        outputFrame(rtp);
    }
    return info.keyFrame();
}

void VP9RtpDecoder::outputFrame(const RtpPacket::Ptr &rtp) {
    if (_frame->dropAble()) {
        // 不参与dts生成  [AUTO-TRANSLATED:dff3b747]
        // Not involved in dts generation
        _frame->_dts = _frame->_pts;
    } else {
        // rtsp没有dts，那么根据pts排序算法生成dts  [AUTO-TRANSLATED:f37c17f3]
        // Rtsp does not have dts, so dts is generated according to the pts sorting algorithm
        _dts_generator.getDts(_frame->_pts, _frame->_dts);
    }

    if (_frame->keyFrame() && _gop_dropped) {
        _gop_dropped = false;
        InfoL << "new gop received, rtp:\r\n" << rtp->dumpString();
    }
    if (!_gop_dropped || _frame->configFrame()) {
        // InfoL << _frame->pts() << " size=" << _frame->size();
        RtpCodec::inputFrame(_frame);
    }
    obtainFrame();
}


////////////////////////////////////////////////////////////////////////

bool VP9RtpEncoder::inputFrame(const Frame::Ptr &frame) {
    uint8_t header[20] = { 0 };
    int nheader = 1;
    header[0] = kBBit;
    bool key = frame->keyFrame();
    if (!key)
        header[0] |= kPBit;
#if 1
    header[0] |= kIBit;
    if (++_pic_id > 0x7FFF) {
        _pic_id = 0;
    }
    header[1] = (0x80 | ((_pic_id >> 8) & 0x7F));
    header[2] = (_pic_id & 0xFF);
    nheader += 2;
#endif
    const char *ptr = frame->data() + frame->prefixSize();
    int len = frame->size() - frame->prefixSize();
    int pdu_size = getRtpInfo().getMaxSize() - nheader;

    bool mark = false;
    for (int pos = 0; pos < len; pos += pdu_size) {
        if (len - pos <= pdu_size) {
            pdu_size = len - pos;
            header[0] |= kEBit;
            mark = true;
        }

        auto rtp = getRtpInfo().makeRtp(TrackVideo, nullptr, pdu_size + nheader, mark, frame->pts());
        if (rtp) {
            uint8_t *payload = rtp->getPayload();
            memcpy(payload, header, nheader);
            memcpy(payload + nheader, ptr + pos, pdu_size);
            RtpCodec::inputRtp(rtp, key);
        }
        key = false;
        header[0] &= (~kBBit); //  Clear 'Begin of partition' bit.
    }
    return true;
}

} // namespace mediakit