ProtobufJsonFormat.scala

import com.fasterxml.jackson.core.Base64Variants
import com.google.protobuf.ByteString
import com.google.protobuf.Descriptors.FieldDescriptor.JavaType
import com.google.protobuf.Descriptors.{ EnumDescriptor, EnumValueDescriptor, FieldDescriptor }
import com.trueaccord.scalapb.{ GeneratedMessage, GeneratedMessageCompanion, Message }
import play.api.data.validation.ValidationError
import play.api.libs.json._

import scala.util.{ Failure, Success, Try }

trait ProtobufJsonFormat {

  def toJson[A](m: GeneratedMessage): JsObject = {
    JsObject(m.getAllFields.map {
      case (fd, v) => fd.getJsonName -> serializeField(fd, v)
    })
  }

  def fromJson[A <: GeneratedMessage with Message[A]](
    value: JsValue
  )(implicit cmp: GeneratedMessageCompanion[A]): JsResult[A] = {
    import scala.collection.JavaConverters._

    def tryParseSingle[B](jsPath: JsPath, toTry: => B): JsResult[B] = Try(toTry) match {
      case Success(a) => JsSuccess(a)
      case Failure(e) => JsError(jsPath -> ValidationError(e.getMessage))
    }

    def parseValue(parentPath: JsPath, fd: FieldDescriptor, value: JsValue): JsResult[Any] =
      if (fd.isRepeated) parseRepeatedValue(parentPath \ fd.getJsonName, fd, value)
      else parseSingleValue(parentPath \ fd.getJsonName, fd, value)

    def parseRepeatedValue(path: JsPath, fd: FieldDescriptor, value: JsValue): JsResult[Seq[Any]] = value match {
      case JsArray(vals) =>
        //Parsing each elements
        val parsedElements = vals.zipWithIndex.map {
          case (jsValue, index) =>
            parseSingleValue(path(index), fd, jsValue)
        }
        // Merging the errors or aggregating all the parsed elements into a sequence
        val sequenced = parsedElements.foldLeft[JsResult[Seq[Any]]](JsSuccess(Seq(), path)) {
          case (JsSuccess(l, _), JsSuccess(r, _)) => JsSuccess(l :+ r, path)
          case (JsError(e1), JsError(e2))         => JsError(JsError.merge(e1, e2))
          case (JsError(e), _)                    => JsError(e)
          case (_, JsError(e))                    => JsError(e)
        }
        sequenced
      case _ =>
        val errorMessage = s"Expected an array for repeated field ${fd.getJsonName} of ${fd.getContainingType.getName}"
        JsError(path -> ValidationError(errorMessage))
    }

    def parseEnum(ed: EnumDescriptor, name: String) = {
      val value = ed.findValueByName(name)
      if (value == null)
        throw new RuntimeException(s"unrecognized enum value: '$name'") // will be catched by tryParseSingle
      else
        value
    }

    def parseSingleValue(path: JsPath, fd: FieldDescriptor, value: JsValue): JsResult[Any] = {

      (fd.getJavaType, value) match {
        case (JavaType.ENUM, JsString(s))    => tryParseSingle(path, parseEnum(fd.getEnumType, s))
        case (JavaType.MESSAGE, o: JsObject) =>
          // The asInstanceOf[] is a lie: we actually have a companion of some other message (not A),
          // but this doesn't matter after erasure.
          fromJson(o)(cmp.messageCompanionForField(fd).asInstanceOf[GeneratedMessageCompanion[A]]).repath(path)
        case (JavaType.INT, JsNumber(x))      => JsSuccess(x.intValue(), path)
        case (JavaType.INT, JsNull)           => JsSuccess(0, path)
        case (JavaType.LONG, JsNumber(x))     => JsSuccess(x.longValue(), path)
        case (JavaType.LONG, JsNull)          => JsSuccess(0L, path)
        case (JavaType.DOUBLE, JsNumber(x))   => JsSuccess(x.doubleValue(), path)
        case (JavaType.DOUBLE, JsNull)        => JsSuccess(0.toDouble, path)
        case (JavaType.FLOAT, JsNumber(x))    => JsSuccess(x.floatValue(), path)
        case (JavaType.FLOAT, JsNull)         => JsSuccess(0.toFloat, path)
        case (JavaType.BOOLEAN, JsBoolean(b)) => JsSuccess(b, path)
        case (JavaType.BOOLEAN, JsNull)       => JsSuccess(false, path)
        case (JavaType.STRING, JsString(s))   => JsSuccess(s, path)
        case (JavaType.STRING, JsNull)        => JsSuccess("", path)
        case (JavaType.BYTE_STRING, JsString(s)) =>
          tryParseSingle(path, ByteString.copyFrom(Base64Variants.getDefaultVariant.decode(s)))
        case (JavaType.BYTE_STRING, JsNull) => JsSuccess(ByteString.EMPTY, path)
        case _ =>
          JsError(
            path -> ValidationError(
              s"Unexpected value ($value) for field ${fd.getJsonName} of ${fd.getContainingType.getName}"
            )
          )
      }
    }

    value match {
      case JsObject(fields) =>
        val values: Map[String, JsValue] = fields.map(k => k._1 -> k._2).toMap

        // Parsing each value of the jsObj independently
        val resultMap: Map[FieldDescriptor, JsResult[Any]] = (for {
          fd      <- cmp.descriptor.getFields.asScala
          jsValue <- values.get(fd.getJsonName)
        } yield (fd, parseValue(JsPath, fd, jsValue))).toMap

        // Merging the errors or aggregating all the parsed elements into a sequence
        val sequenced = resultMap.foldLeft[JsResult[Map[FieldDescriptor, Any]]](JsSuccess(Map())) {
          case (JsSuccess(acc, _), (fd, JsSuccess(v, _))) => JsSuccess(acc + (fd -> v))
          case (JsError(e1), (_, JsError(e2)))            => JsError(JsError.merge(e1, e2))
          case (JsError(e), _)                            => JsError(e)
          case (_, (_, JsError(e)))                       => JsError(e)
        }

        // The fromFieldsMap operation might still fail.
        sequenced.flatMap(valueMap => tryParseSingle(JsPath, cmp.fromFieldsMap(valueMap)))
      case _ => JsError(s"Expected an object, found $value")
    }
  }

  @inline
  private def serializeField(fd: FieldDescriptor, value: Any): JsValue = {
    if (fd.isRepeated) {
      JsArray(value.asInstanceOf[Seq[Any]].map(serializeSingleValue(fd, _)).toList)
    } else serializeSingleValue(fd, value)
  }

  @inline
  private def serializeSingleValue(fd: FieldDescriptor, value: Any): JsValue = fd.getJavaType match {
    case JavaType.ENUM    => JsString(value.asInstanceOf[EnumValueDescriptor].getName)
    case JavaType.MESSAGE => toJson(value.asInstanceOf[GeneratedMessage])
    case JavaType.INT     => JsNumber(BigDecimal(value.asInstanceOf[Int]))
    case JavaType.LONG    => JsNumber(BigDecimal(value.asInstanceOf[Long]))
    case JavaType.DOUBLE  => JsNumber(BigDecimal(value.asInstanceOf[Double]))
    case JavaType.FLOAT   => JsNumber(BigDecimal.decimal(value.asInstanceOf[Float]))
    case JavaType.BOOLEAN => JsBoolean(value.asInstanceOf[Boolean])
    case JavaType.STRING  => JsString(value.asInstanceOf[String])
    case JavaType.BYTE_STRING =>
      JsString(Base64Variants.getDefaultVariant.encode(value.asInstanceOf[ByteString].toByteArray))
  }

  implicit def protoToReader[T <: GeneratedMessage with Message[T]: GeneratedMessageCompanion]: Reads[T] =
    new Reads[T] {
      def reads(value: JsValue): JsResult[T] = fromJson(value)
    }

  implicit def protoToWriter[T <: GeneratedMessage with Message[T]: GeneratedMessageCompanion]: Writes[T] =
    new Writes[T] {
      def writes(obj: T): JsValue = toJson(obj)
    }
}

object ProtobufJsonFormat extends ProtobufJsonFormat