/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * license agreements; and to You under the Apache License, version 2.0:
 *
 *   https://www.apache.org/licenses/LICENSE-2.0
 *
 * This file is part of the Apache Pekko project, which was derived from Akka.
 */

/*
 * Copyright (C) 2018-2021 Lightbend Inc. <https://www.lightbend.com>
 */

package org.apache.pekko.grpc.internal

import org.apache.pekko
import pekko.NotUsed
import pekko.actor.{ ActorSystem, ClassicActorSystemProvider }
import pekko.annotation.InternalApi
import pekko.grpc.GrpcProtocol.{ GrpcProtocolWriter, TrailerFrame }
import pekko.grpc.scaladsl.{ headers, GrpcExceptionHandler }
import pekko.grpc.{ ProtobufFrameSerializer, ProtobufSerializer, Trailers }
import pekko.http.scaladsl.model.HttpEntity.ChunkStreamPart
import pekko.http.scaladsl.model.{
  AttributeKey,
  AttributeKeys,
  HttpEntity,
  HttpHeader,
  HttpProtocols,
  HttpResponse,
  StatusCodes,
  Trailer
}
import pekko.stream.Materializer
import pekko.stream.scaladsl.Source
import pekko.util.ByteString
import io.grpc.Status

import scala.collection.immutable
import scala.concurrent.{ ExecutionContext, Future }
import scala.util.control.NonFatal

/**
 * Some helpers for creating HTTP entities for use with gRPC.
 *
 * INTERNAL API
 */
@InternalApi // consumed from generated classes so cannot be private
object GrpcResponseHelpers {
  private val TrailerOk = GrpcEntityHelpers.trailer(Status.OK)
  private val TrailerOkAttribute = Trailer(TrailerOk.trailers)
  private val TrailerOkAttributes =
    Map.empty[AttributeKey[_], Any].updated(AttributeKeys.trailer, TrailerOkAttribute)
  private val IdentityResponseHeaders: immutable.Seq[HttpHeader] =
    headers.`Message-Encoding`(Identity.name) :: Nil

  def apply[T](e: Source[T, NotUsed])(
      implicit m: ProtobufSerializer[T],
      writer: GrpcProtocolWriter,
      system: ClassicActorSystemProvider): HttpResponse =
    GrpcResponseHelpers(e, Source.single(TrailerOk))

  def apply[T](e: Source[T, NotUsed], eHandler: ActorSystem => PartialFunction[Throwable, Trailers])(
      implicit m: ProtobufSerializer[T],
      writer: GrpcProtocolWriter,
      system: ClassicActorSystemProvider): HttpResponse =
    GrpcResponseHelpers(e, Source.single(TrailerOk), eHandler)

  def responseForSingleElement[T](e: T, eHandler: ActorSystem => PartialFunction[Throwable, Trailers])(
      implicit m: ProtobufSerializer[T],
      writer: GrpcProtocolWriter,
      system: ClassicActorSystemProvider): HttpResponse = {
    val responseHeaders = responseHeadersFor(writer)
    try {
      if ((writer.messageEncoding eq Identity) && writer.contentType == GrpcProtocolNative.contentType) {
        m match {
          case frameSerializer: ProtobufFrameSerializer[T @unchecked] =>
            nativeResponse(writer, frameSerializer.serializeDataFrame(e), responseHeaders)
          case _ =>
            writer.encodeDataToResponse(m.serialize(e), responseHeaders, TrailerOkAttribute)
        }
      } else {
        writer.encodeDataToResponse(m.serialize(e), responseHeaders, TrailerOkAttribute)
      }
    } catch {
      case NonFatal(ex) => status(GrpcEntityHelpers.handleException(ex, eHandler))
    }
  }

  private def responseHeadersFor(writer: GrpcProtocolWriter): immutable.Seq[HttpHeader] =
    if (writer.messageEncoding eq Identity) IdentityResponseHeaders
    else headers.`Message-Encoding`(writer.messageEncoding.name) :: Nil

  private def nativeResponse(
      writer: GrpcProtocolWriter,
      encodedData: ByteString,
      responseHeaders: immutable.Seq[HttpHeader]): HttpResponse =
    new HttpResponse(
      status = StatusCodes.OK,
      headers = responseHeaders,
      entity = HttpEntity(writer.contentType, encodedData),
      protocol = HttpProtocols.`HTTP/1.1`,
      attributes = TrailerOkAttributes)

  def apply[T](e: Source[T, NotUsed], status: Future[Status])(
      implicit m: ProtobufSerializer[T],
      mat: Materializer,
      writer: GrpcProtocolWriter,
      system: ClassicActorSystemProvider): HttpResponse =
    GrpcResponseHelpers(e, status, GrpcExceptionHandler.defaultMapper _)

  def apply[T](
      e: Source[T, NotUsed],
      status: Future[Status],
      eHandler: ActorSystem => PartialFunction[Throwable, Trailers])(
      implicit m: ProtobufSerializer[T],
      mat: Materializer,
      writer: GrpcProtocolWriter,
      system: ClassicActorSystemProvider): HttpResponse = {
    implicit val ec: ExecutionContext = mat.executionContext
    GrpcResponseHelpers(
      e,
      Source.lazyFuture(() => status.map(GrpcEntityHelpers.trailer(_))).mapMaterializedValue(_ => NotUsed),
      eHandler)
  }

  def apply[T](
      e: Source[T, NotUsed],
      trail: Source[TrailerFrame, NotUsed],
      eHandler: ActorSystem => PartialFunction[Throwable, Trailers] = GrpcExceptionHandler.defaultMapper)(
      implicit m: ProtobufSerializer[T],
      writer: GrpcProtocolWriter,
      system: ClassicActorSystemProvider): HttpResponse = {
    response(GrpcEntityHelpers(e, trail, eHandler))
  }

  private def response(entity: Source[ChunkStreamPart, NotUsed])(implicit writer: GrpcProtocolWriter) = {
    HttpResponse(
      headers = immutable.Seq(
        headers.`Message-Encoding`(writer.messageEncoding.name),
        // Pre-announcing trailers: See https://www.rfc-editor.org/rfc/rfc7230 #Section 4.4
        headers.`Trailer`(headers.`Status`.name)
      ),
      entity = HttpEntity.Chunked(writer.contentType, entity))
  }

  def status(trailer: Trailers)(implicit writer: GrpcProtocolWriter): HttpResponse = {
    HttpResponse(
      headers =
        headers.`Message-Encoding`(writer.messageEncoding.name) ::
        GrpcEntityHelpers.trailers(trailer.status, trailer.metadata),
      entity = HttpEntity.empty(writer.contentType)
    )
  }
}
