_ext.cpp

#include <pybind11/pybind11.h>
#include <pybind11/stl.h>

#include <cstdint>
#include <cstring>
#include <limits>
#include <stdexcept>
#include <string>
#include <string_view>
#include <vector>

namespace py = pybind11;

namespace {

class BufferView {
 public:
  explicit BufferView(py::handle object) {
    if (PyObject_GetBuffer(object.ptr(), &view_, PyBUF_CONTIG_RO) != 0) {
      throw py::type_error(
          "input must be a contiguous bytes-like object");
    }
    if (view_.len < 0) {
      PyBuffer_Release(&view_);
      throw py::value_error("input buffer has a negative length");
    }
    acquired_ = true;
  }

  BufferView(const BufferView&) = delete;
  BufferView& operator=(const BufferView&) = delete;

  ~BufferView() {
    if (acquired_) {
      PyBuffer_Release(&view_);
    }
  }

  const uint8_t* data() const {
    return static_cast<const uint8_t*>(view_.buf);
  }

  size_t size() const {
    return static_cast<size_t>(view_.len);
  }

 private:
  Py_buffer view_{};
  bool acquired_ = false;
};

struct Vp8Frame {
  const uint8_t* data = nullptr;
  size_t size = 0;
  uint16_t width = 0;
  uint16_t height = 0;
};

uint16_t read_le16(const uint8_t* data) {
  return static_cast<uint16_t>(data[0]) |
      (static_cast<uint16_t>(data[1]) << 8);
}

uint32_t read_le32(const uint8_t* data) {
  return static_cast<uint32_t>(data[0]) |
      (static_cast<uint32_t>(data[1]) << 8) |
      (static_cast<uint32_t>(data[2]) << 16) |
      (static_cast<uint32_t>(data[3]) << 24);
}

bool tag_equals(const uint8_t* tag, const char (&expected)[5]) {
  return std::memcmp(tag, expected, 4) == 0;
}

[[noreturn]] void unsupported(const std::string& message) {
  throw py::value_error("unsupported WebP: " + message);
}

Vp8Frame validate_vp8_payload(const uint8_t* data, size_t size) {
  if (size < 10) {
    throw py::value_error("invalid VP8 payload: key-frame header is truncated");
  }
  if ((data[0] & 1U) != 0) {
    throw py::value_error(
        "invalid VP8 payload: WebP must contain a VP8 key frame");
  }
  if (data[3] != 0x9d || data[4] != 0x01 || data[5] != 0x2a) {
    throw py::value_error("invalid VP8 payload: key-frame start code is missing");
  }

  const uint16_t width = read_le16(data + 6) & 0x3fffU;
  const uint16_t height = read_le16(data + 8) & 0x3fffU;
  if (width == 0 || height == 0) {
    throw py::value_error("invalid VP8 payload: dimensions must be non-zero");
  }

  return Vp8Frame{data, size, width, height};
}

Vp8Frame parse_webp(const uint8_t* data, size_t size) {
  if (size < 12) {
    throw py::value_error("invalid WebP: RIFF header is truncated");
  }
  if (!tag_equals(data, "RIFF")) {
    throw py::value_error("invalid WebP: missing RIFF signature");
  }
  if (!tag_equals(data + 8, "WEBP")) {
    throw py::value_error("invalid WebP: missing WEBP signature");
  }

  const uint32_t riff_size = read_le32(data + 4);
  if (riff_size < 4) {
    throw py::value_error("invalid WebP: RIFF size is smaller than WEBP form");
  }
  const uint64_t riff_end64 = 8ULL + static_cast<uint64_t>(riff_size);
  if (riff_end64 > size) {
    throw py::value_error("invalid WebP: RIFF payload is truncated");
  }
  const size_t riff_end = static_cast<size_t>(riff_end64);

  Vp8Frame result;
  bool found_vp8 = false;
  size_t position = 12;
  while (position < riff_end) {
    if (riff_end - position < 8) {
      throw py::value_error("invalid WebP: RIFF chunk header is truncated");
    }

    const uint8_t* tag = data + position;
    const uint32_t chunk_size = read_le32(tag + 4);
    const uint64_t payload_start = static_cast<uint64_t>(position) + 8;
    const uint64_t payload_end =
        payload_start + static_cast<uint64_t>(chunk_size);
    const uint64_t padded_end = payload_end + (chunk_size & 1U);
    if (payload_end > riff_end || padded_end > riff_end) {
      throw py::value_error("invalid WebP: RIFF chunk payload is truncated");
    }

    const uint8_t* payload = data + static_cast<size_t>(payload_start);
    if (tag_equals(tag, "VP8 ")) {
      if (found_vp8) {
        throw py::value_error("invalid WebP: multiple VP8 chunks");
      }
      result = validate_vp8_payload(payload, chunk_size);
      found_vp8 = true;
    } else if (tag_equals(tag, "VP8L")) {
      unsupported("VP8L lossless bitstreams cannot be decoded by NVDEC");
    } else if (tag_equals(tag, "ALPH")) {
      unsupported("embedded ALPH chunks require a separate alpha decoder");
    } else if (tag_equals(tag, "ANIM") || tag_equals(tag, "ANMF")) {
      unsupported("animated WebP is not supported");
    }

    position = static_cast<size_t>(padded_end);
  }

  if (!found_vp8) {
    throw py::value_error("invalid WebP: no lossy VP8 chunk found");
  }
  return result;
}

void append_le16(std::string& output, uint16_t value) {
  output.push_back(static_cast<char>(value & 0xffU));
  output.push_back(static_cast<char>((value >> 8) & 0xffU));
}

void append_le32(std::string& output, uint32_t value) {
  for (int shift = 0; shift < 32; shift += 8) {
    output.push_back(static_cast<char>((value >> shift) & 0xffU));
  }
}

void append_le64(std::string& output, uint64_t value) {
  for (int shift = 0; shift < 64; shift += 8) {
    output.push_back(static_cast<char>((value >> shift) & 0xffU));
  }
}

void validate_ivf_arguments(
    size_t frame_count,
    uint16_t width,
    uint16_t height,
    uint32_t rate,
    uint32_t scale) {
  if (frame_count == 0) {
    throw py::value_error("at least one VP8 frame is required");
  }
  if (frame_count > std::numeric_limits<uint32_t>::max()) {
    throw std::overflow_error("IVF frame count exceeds uint32");
  }
  if (width == 0 || height == 0) {
    throw py::value_error("IVF dimensions must be non-zero");
  }
  if (rate == 0 || scale == 0) {
    throw py::value_error("IVF rate and scale must be non-zero");
  }
}

std::string make_ivf(
    const std::vector<Vp8Frame>& frames,
    uint16_t width,
    uint16_t height,
    uint32_t rate,
    uint32_t scale) {
  validate_ivf_arguments(frames.size(), width, height, rate, scale);

  uint64_t output_size = 32;
  for (const Vp8Frame& frame : frames) {
    if (frame.width != width || frame.height != height) {
      throw py::value_error(
          "all VP8 frames in an IVF stream must have identical dimensions");
    }
    if (frame.size > std::numeric_limits<uint32_t>::max()) {
      throw std::overflow_error("VP8 frame size exceeds IVF uint32 field");
    }
    output_size += 12ULL + frame.size;
    if (output_size >
        static_cast<uint64_t>(std::numeric_limits<Py_ssize_t>::max())) {
      throw std::overflow_error(
          "IVF output is too large for a Python bytes object");
    }
  }

  std::string output;
  output.reserve(static_cast<size_t>(output_size));
  output.append("DKIF", 4);
  append_le16(output, 0);
  append_le16(output, 32);
  output.append("VP80", 4);
  append_le16(output, width);
  append_le16(output, height);
  append_le32(output, rate);
  append_le32(output, scale);
  append_le32(output, static_cast<uint32_t>(frames.size()));
  append_le32(output, 0);

  for (size_t index = 0; index < frames.size(); ++index) {
    const Vp8Frame& frame = frames[index];
    append_le32(output, static_cast<uint32_t>(frame.size));
    append_le64(output, static_cast<uint64_t>(index));
    output.append(
        reinterpret_cast<const char*>(frame.data),
        static_cast<std::string::size_type>(frame.size));
  }
  return output;
}

py::tuple extract_vp8(py::handle webp) {
  BufferView input(webp);
  const Vp8Frame frame = parse_webp(input.data(), input.size());
  return py::make_tuple(
      py::bytes(reinterpret_cast<const char*>(frame.data), frame.size),
      frame.width,
      frame.height);
}

py::bytes build_ivf(
    py::iterable vp8_payloads,
    uint16_t width,
    uint16_t height,
    uint32_t rate,
    uint32_t scale) {
  std::vector<std::string> owned_payloads;
  for (py::handle payload : vp8_payloads) {
    BufferView input(payload);
    owned_payloads.emplace_back(
        reinterpret_cast<const char*>(input.data()), input.size());
  }

  std::vector<Vp8Frame> frames;
  frames.reserve(owned_payloads.size());
  for (const std::string& payload : owned_payloads) {
    frames.push_back(validate_vp8_payload(
        reinterpret_cast<const uint8_t*>(payload.data()), payload.size()));
  }

  const std::string output = make_ivf(frames, width, height, rate, scale);
  return py::bytes(output);
}

py::bytes webp_batch_to_ivf(
    py::iterable webps,
    uint32_t rate,
    uint32_t scale) {
  std::vector<std::string> owned_payloads;
  uint16_t width = 0;
  uint16_t height = 0;

  for (py::handle webp : webps) {
    BufferView input(webp);
    const Vp8Frame frame = parse_webp(input.data(), input.size());
    if (owned_payloads.empty()) {
      width = frame.width;
      height = frame.height;
    } else if (frame.width != width || frame.height != height) {
      throw py::value_error(
          "all WebPs in an IVF stream must have identical dimensions");
    }
    owned_payloads.emplace_back(
        reinterpret_cast<const char*>(frame.data), frame.size);
  }

  std::vector<Vp8Frame> frames;
  frames.reserve(owned_payloads.size());
  for (const std::string& payload : owned_payloads) {
    frames.push_back(validate_vp8_payload(
        reinterpret_cast<const uint8_t*>(payload.data()), payload.size()));
  }

  const std::string output = make_ivf(frames, width, height, rate, scale);
  return py::bytes(output);
}

}  // namespace

PYBIND11_MODULE(_vp8_ivf_ext, module) {
  module.doc() = "Strict lossy-WebP VP8 parser and in-memory IVF builder";
  module.def(
      "extract_vp8",
      &extract_vp8,
      py::arg("webp"),
      "Return (VP8 payload, width, height) from a supported lossy WebP.");
  module.def(
      "build_ivf",
      &build_ivf,
      py::arg("vp8_payloads"),
      py::arg("width"),
      py::arg("height"),
      py::arg("rate") = 25,
      py::arg("scale") = 1,
      "Build an IVF byte stream from same-resolution VP8 key frames.");
  module.def(
      "webp_batch_to_ivf",
      &webp_batch_to_ivf,
      py::arg("webps"),
      py::arg("rate") = 25,
      py::arg("scale") = 1,
      "Extract same-resolution lossy WebPs and return one IVF byte stream.");
}