# Copyright 2014-2020 Chris Cummins <chrisc.101@gmail.com>.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Utility code for working with Protocol Buffers."""
import collections
import gzip
import json
import pathlib
import subprocess
import typing

import google.protobuf.json_format
import google.protobuf.message
import google.protobuf.text_format

# A type alias for annotating methods which take or return protocol buffers.
ProtocolBuffer = typing.Any

# A type alias for protocol enum fields.
Enum = int


class ProtoValueError(ValueError):
  """Raised in case of a value error from a proto."""

  pass


class EncodeError(ProtoValueError):
  """Raised in case of error encoding a proto."""

  pass


class DecodeError(ProtoValueError):
  """Raised in case of error decoding a proto."""

  pass


class ProtoWorkerTimeoutError(subprocess.CalledProcessError):
  """Raised is a protobuf worker binary times out."""

  def __init__(
    self, cmd: typing.List[str], timeout_seconds: int, returncode: int,
  ):
    self.cmd = cmd
    self.timeout_seconds = t
imeout_seconds
    # subprocess.CalledProcessError.str() requires a ret
urncode attribute.
    self.returncode = returncode

  def __repr__(self) -> str:
    return (
      f"Proto worker timeout after {self.timeout_seconds} "
      f"seconds: {' '.join(self.cmd)}"
    )


def FromString(
  string: str, message: ProtocolBuffer, uninitialized_okay: bool = False,
) -> ProtocolBuffer:
  """Read a text format protocol buffer from a string.

  Args:
    string: A text format protocol buffer.
    message: A message instance to read into.
    uninitialized_okay: If True, do not require that decoded messages be
      initialized. If False, DecodeError is raised.

  Returns:
    The parsed message (same as the message argument).

  Raises:
    DecodeError: If the file cannot be decoded to the given message type, or if
      after decoding, the message is not initialized and uninitialized_okay is
      False.
  """
  try:
    google.protobuf.text_format.Merge(string, message)
  except google.protobuf.text_format.ParseError as e:
    raise DecodeError(e)

  if not uninitialized_okay and not message.IsInitialized():
    raise DecodeError(f"Required fields not set")

  return message


def FromFile(
  path: pathlib.Path,
  message: ProtocolBuffer,
  assume_filename: typing.Optional[typing.Union[str, pathlib.Path]] = None,
  uninitialized_okay: bool = False,
) -> ProtocolBuffer:
  """Read a protocol buffer from a file.

  This method uses attempts to guess the encoding from the path suffix,
  supporting binary, text, and json formatted messages. The mapping of suffixes
  to formatting is, in order:
      *.txt.gz: Gzipped text.
      *.txt: Text.
      *.pbtxt.gz: Gzipped text.
      *.pbtxt: Text.
      *.json.gz: Gzipped JSON.
      *.json: JSON.
      *.gz: Gzipped encoded string.
      *: Encoded string.

  Args:
    path: Path to the proto file.
    message: A message instance to read into.
    assume_filename: For the purpose of determining the encoding from the file
      extension, use this name rather than the true path.
    uninitialized_okay: If True, do not require that decoded messages be
      initialized. If False, DecodeError is raised.

  Returns:
    The parsed message (same as the message argument).

  Raises:
    FileNotFoundError: If the path does not exist.
    IsADirectoryError: If the path is a directory.
    DecodeError: If the file cannot be decoded to the given message type, or if
      after decoding, the message is not initialized and uninitialized_okay is
      False.
  """
  if not path.is_file():
    if path.is_dir():
      raise IsADirectoryError(f"Path is a directory: '{path}'")
    else:
      raise FileNotFoundError(f"File not found: '{path}'")

  suffixes = (
    pathlib.Path(assume_filename,).suffixes
    if assume_filename
    else path.suffixes
  )
  if suffixes and suffixes[-1] == ".gz":
    suffixes.pop()
    open_function = gzip.open
  else:
    open_function = open

  suffix = suffixes[-1] if suffixes else ""
  try:
    with open_function(path, "rb") as f:
      if suffix == ".txt" or suffix == ".pbtxt":
        # Allow uninitialized fields here because we will catch the error later,
        # allowing us to report the path of the proto.
        FromString(f.read().decode("utf-8"), message, uninitialized_okay=True)
      elif suffix == ".json":
        google.protobuf.json_format.Parse(f.read(), message)
      else:
        message.ParseFromString(f.read())
  except (
    google.protobuf.message.DecodeError,
    google.protobuf.text_format.ParseError,
    google.protobuf.json_format.ParseError,
  ) as e:
    # The exception raised during parsing depends on the message format. Catch
    # them all under a single DecodeError exception type.
    raise DecodeError(e)

  if not uninitialized_okay and not message.IsInitialized():
    raise DecodeError(f"Required fields not set: '{path}'")

  return message


def ToFile(
  message: ProtocolBuffer,
  path: pathlib.Path,
  exist_ok: bool = True,
  assume_filename: typing.Optional[typing.Union[str, pathlib.Path]] = None,
) -> ProtocolBuffer:
  """Write a protocol buffer to a file.

  This method uses attempts to guess the encoding from the path suffix,
  supporting binary, text, and json formatted messages. The mapping of suffixes
  to formatting is, in order:
      *.txt.gz: Gzipped text format.
      *.txt: Text format.
      *.pbtxt.gz: Gzipped text format.
      *.pbtxt: Text format.
      *.json.gz: Gzipped JSON format.
      *.json: JSON format.
      *.gz: Gzipped binary format.
      *: Binary format.

  Args:
    message: A message instance to write to file. The message must be
      initialized, i.e. have all required fields set.
    path: Path to the proto file.
    exist_ok: If True, overwrite existing file.
    assume_filename: For the purpose of determining the encoding from the file
      extension, use this name rather than the true path.

  Returns:
    The parsed message (same as the message argument).

  Raises:
    EncodeError: If the message is not initialized, i.e. it is missing required
      fields.
    FileNotFoundError: If the parent directory of the requested path does not
      exist.
    IsADirectoryError: If the requested path is a directory.
    FileExistsError: If the requested path already exists and exist_ok is False.
  """
  if not exist_ok and path.exists():
    raise FileExistsError(f"Refusing to overwrite {path}")

  # The SerializeToString() method refuses to encode a message which is not
  # initialized, whereas the MessageToString() and MessageToJson() methods do
  # not. This API should be consistent, so we enforce that all formats require
  # the message to be initialized.
  if not message.IsInitialized():
    class_name = type(message).__name__
    raise EncodeError(f"Required fields not set: '{class_name}'")

  suffixes = (
    pathlib.Path(assume_filename,).suffixes
    if assume_filename
    else path.suffixes
  )
  if suffixes and suffixes[-1] == ".gz":
    suffixes.pop()
    open_function = gzip.open
  else:
    open_function = open

  suffix = suffixes[-1] if suffixes else ""
  mode = "wt" if suffix in {".txt", ".pbtxt", ".json"} else "wb"

  with open_function(path, mode) as f:
    if suffix == ".txt" or suffix == ".pbtxt":
      f.write(google.protobuf.text_format.MessageToString(message))
    elif suffix == ".json":
      f.write(
        google.protobuf.json_format.MessageToJson(
          message, preserving_proto_field_name=True,
        ),
      )
    else:
      f.write(message.SerializeToString())

  return message


def ToJson(message: ProtocolBuffer) -> "jsonutil.JSO