/*
 * 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.
 */

package org.apache.pekko.grpc.interop

import org.apache.pekko
import pekko.stream.{ Materializer, SystemMaterializer }
import pekko.http.scaladsl.model._
import pekko.http.scaladsl.model.headers.RawHeader
import pekko.http.scaladsl.server.RouteResult.Complete
import pekko.http.scaladsl.server.{ Directive0, Directives, Route }

import io.grpc.testing.integration.TestServiceHandlerFactory
import io.grpc.testing.integration.test.TestService

import scala.collection.immutable
import scala.concurrent.ExecutionContext

// Generated by our plugin
import io.grpc.testing.integration.test.TestServiceHandler

class GrpcInteropIoWithPekkoSpec extends GrpcInteropTests(IoGrpcJavaServerProvider, PekkoHttpClientProviderScala)

class GrpcInteropPekkoScalaWithIoSpec extends GrpcInteropTests(PekkoHttpServerProviderScala, IoGrpcJavaClientProvider)
class GrpcInteropPekkoScalaWithPekkoSpec
    extends GrpcInteropTests(PekkoHttpServerProviderScala, PekkoHttpClientProviderScala)

class GrpcInteropPekkoJavaWithIoSpec extends GrpcInteropTests(PekkoHttpServerProviderJava, IoGrpcJavaClientProvider)
class GrpcInteropPekkoJavaWithPekkoSpec
    extends GrpcInteropTests(PekkoHttpServerProviderJava, PekkoHttpClientProviderScala)

object PekkoHttpServerProviderScala extends PekkoHttpServerProvider {
  import Directives._

  val label: String = "pekko-grpc server scala"
  val pendingCases = Set()

  val server = PekkoGrpcServerScala(implicit sys => {
    implicit val ec: ExecutionContext = sys.dispatcher

    val requestHandler = TestServiceHandler(new TestServiceImpl())

    val route: Route = (pathPrefix(TestService.name) & echoHeaders) { ctx =>
      requestHandler(ctx.request).map(Complete)
    }

    Route.toFunction(Route.seal(route))
  })

  val echoHeaders: Directive0 = extractRequest.flatMap(request => {
    val initialHeaderToEcho = request.headers.find(_.name() == "x-grpc-test-echo-initial")
    val trailingHeaderToEcho = request.headers.find(_.name() == "x-grpc-test-echo-trailing-bin")

    mapResponseHeaders(h => h ++ initialHeaderToEcho) & mapTrailingResponseHeaders(h => h ++ trailingHeaderToEcho)
  })

  // TODO to be moved to the runtime lib (or even pekko-http itself?)
  def mapTrailingResponseHeaders(f: immutable.Seq[HttpHeader] => immutable.Seq[HttpHeader]): Directive0 =
    mapResponse(response =>
      response.entity match {
        case HttpEntity.Chunked(contentType, data) =>
          response.withEntity(
            HttpEntity.Chunked(
              contentType,
              data.map {
                case chunk: HttpEntity.Chunk    => chunk
                case last: HttpEntity.LastChunk =>
                  HttpEntity.LastChunk(last.extension, f(last.trailer))
              }))
        case _ =>
          response
            .attribute(AttributeKeys.trailer)
            .map(trailer => f(trailer.headers.map((RawHeader.apply _).tupled)))
            .flatMap(headers => if (headers.isEmpty) None else Some(Trailer(headers)))
            .fold(response)(response.addAttribute(AttributeKeys.trailer, _))
      })
}

object PekkoHttpServerProviderJava extends PekkoHttpServerProvider {
  val label: String = "pekko-grpc server java"

  val pendingCases =
    Set(
      "custom_metadata")

  val server = new PekkoGrpcServerJava((mat, sys) => {
    TestServiceHandlerFactory.create(new JavaTestServiceImpl(mat), sys)
  })
}

object PekkoHttpClientProviderScala extends PekkoClientProvider {
  val label: String = "pekko-grpc scala client tester"

  def client = PekkoGrpcClientScala(settings => implicit sys => new PekkoGrpcClientTester(settings))
}
