""" OggOpus parsing code. """

import logging
import struct

import crcmod


# https://tools.ietf.org/html/rfc3533#section-6
# http://www.onicos.com/staff/iz/formats/ogg.html
# https://en.wikipedia.org/wiki/Ogg#Page_structure
OGG_FIRST_PAGE_HEADER = struct.Struct("<4sBBQLLLB")
OGG_FIRST_PAGE_HEADER_CRC_OFFSET = 22
OGG_FIRST_PAGE_HEADER_CRC = struct.Struct("<L")

# https://tools.ietf.org/html/rfc7845#section-5.1
OGG_OPUS_ID_HEADER = struct.Struct("<8sBBHLhB")
OGG_OPUS_ID_HEADER_GAIN_OFFSET = 16
OGG_OPUS_ID_HEADER_GAIN = struct.Struct("<h")

ogg_page_crc = crcmod.mkCrcFun(0x104c11db7, initCrc=0, rev=False)


def logger():
  return logging.getLogger("r128gain.oggopus")


def parse_oggopus_output_gain(file):
  """ Parse an OggOpus file headers, read and return its output gain, and set file seek position to start of Opus header. """
  #
  # Ogg header
  #

  # check fields of Ogg page header
  chunk = file.read(OGG_FIRST_PAGE_HEADER.size)
  first_ogg_page = bytearray()
  first_ogg_page.extend(chunk)
  if len(chunk) < OGG_FIRST_PAGE_HEADER.size:
    logger().error("Not enough bytes in Ogg page header: %u, expected at least %u" % (len(chunk),
                                                                                      OGG_FIRST_PAGE_HEADER.size))
    return
  capture_pattern, version, header_type, granule_position, bitstream_serial_number, page_sequence_number, \
      crc_checksum, page_segments = OGG_FIRST_PAGE_HEADER.unpack(chunk)
  if capture_pattern != b"OggS":
    logger().error("Invalid OGG capture pattern: %s, expected '%s'" % (repr(capture_pattern), "OggS"))
    return
  if version != 0:
    logger().error("Invalid OGG version: %u, expected %u" % (version, 0))
    return
  if header_type != 2:  # should be first page of stream
    logger().error("Invalid OGG page header type: %u, expected %u" % (header_type, 2))
    return
  if page_sequence_number != 0:
    logger().error("Invalid OGG page sequence number: %u, expected %u" % (page_sequence_number, 0))
    return
  segment_table_fmt = struct.Struct("<%uB" % (page_segments))
  chunk = file.read(segment_table_fmt.size)
  first_ogg_page.extend(chunk)
  if len(chunk) < segment_table_fmt.size:
    logger().error("Not enough bytes for OGG segment table: %u, expected at least %u" % (len(chunk),
                                                                                         segment_table_fmt.size))
    return
  segment_table = segment_table_fmt.unpack(chunk)

  # check crc of first page
  first_ogg_page_size = OGG_FIRST_PAGE_HEADER.size + segment_table_fmt.size + sum(segment_table)
  chunk = file.read(sum(segment_table))
  first_ogg_page.extend(chunk)
  if len(first_ogg_page) < first_ogg_page_size:
    logger().error("Not enough bytes for first OGG page: %u, expected at least %u" % (len(first_ogg_page),
                                                                                      first_ogg_page_size))
    return
  computed_crc = _compute_ogg_page_crc(first_ogg_page)
  if computed_crc != crc_checksum:
    logger().error("Invalid OGG page CRC: 0x%08x, expected 0x%08x" % (crc_checksum, computed_crc))
    return

  #
  # Opus header
  #
  chunk = first_ogg_page[OGG_FIRST_PAGE_HEADER.size + segment_table_fmt.size:][:segment_table[0]]
  if len(chunk) < OGG_OPUS_ID_HEADER.size:
    logger().error("Not enough bytes for Opus Identification header: %u, "
                   "expected at least %u" % (len(chunk),
                                             OGG_OPUS_ID_HEADER.size))
    return
  magic, version, channel_count, preskip, input_samplerate, output_gain, \
      mapping_family = OGG_OPUS_ID_HEADER.unpack(chunk[:OGG_OPUS_ID_HEADER.size])
  if magic != b"OpusHead":
    logger().error("Invalid Opus magic number: %s, expected '%s'" % (repr(magic), "OpusHead"))
    return
  if (version >> 4) != 0:
    logger().error("Invalid Opus version: 0x%x, expected 0x0-0xf" % (version))
    return

  # seek to Opus header
  file.seek(OGG_FIRST_PAGE_HEADER.size + segment_table_fmt.size)

  return output_gain


def write_oggopus_output_gain(file, new_output_gain):
  """ Write output gain Opus header for a file.

  file must be an object successfully used by parse_oggopus_output_gain.
  """
  opus_header_pos = file.tell()

  # write Opus header with new gain
  file.seek(opus_header_pos + OGG_OPUS_ID_HEADER_GAIN_OFFSET)
  file.write(OGG_OPUS_ID_HEADER_GAIN.pack(new_output_gain))

  # compute page crc
  file.seek(0)
  page = file.read(opus_header_pos + OGG_OPUS_ID_HEADER.size)
  computed_crc = _compute_ogg_page_crc(page)

  # write CRC
  file.seek(OGG_FIRST_PAGE_HEADER_CRC_OFFSET)
  file.write(OGG_FIRST_PAGE_HEADER_CRC.pack(computed_crc))


def _compute_ogg_page_crc(page):
  """ Compute CRC of an Ogg page. """
  page_zero_crc = page[:OGG_FIRST_PAGE_HEADER_CRC_OFFSET] + \
      b"\00" * OGG_FIRST_PAGE_HEADER_CRC.size + \
      page[OGG_FIRST_PAGE_HEADER_CRC_OFFSET + OGG_FIRST_PAGE_HEADER_CRC.size:]
  return ogg_page_crc(page_zero_crc)
