//设置package包名称以及导入依赖的类
package websocket

import akka.NotUsed
import akka.actor._
import akka.http.scaladsl.model.ws.{BinaryMessage, Message, TextMessage}
import akka.stream.scaladsl._
import akka.stream.{ActorMaterializer, OverflowStrategy}
import core.entities.AuthTokenContext
import websocket.ClientEndpoint._
import websocket.WebSocketHandler._

import scala.concurrent.ExecutionContext
import scala.concurrent.duration.FiniteDuration

class WebSocketHandler(processingRouter: ActorRef, connectedClientsStore: ActorRef, processingTimeout: FiniteDuration)
                      (implicit system: ActorSystem, mat: ActorMaterializer, apiDispatcher: ExecutionContext) {

  def clientEndpoint(ctx: AuthTokenContext): Flow[Message, Message, NotUsed] = {
    val clientEndpoint =
      system.actorOf(ClientEndpoint.props(ctx, processingRouter, connectedClientsStore, processingTimeout),
                     ClientEndpoint.name(ctx))

    val incomingMessages: Sink[Message, NotUsed] = Flow[Message]
      .map {
        case TextMessage.Strict(jsContent) => Some(IncomingMessage(jsContent))

        case ts: TextMessage.Streamed =>
          ts.textStream.runWith(Sink.ignore)
          None

        case br: BinaryMessage =>
          br.dataStream.runWith(Sink.ignore)
          None
      }
      .collect {
        case Some(message: IncomingMessage) => message
      }
      .to(Sink.actorRef[IncomingMessage](clientEndpoint, PoisonPill))

    val outgoingMessages: Source[Message, NotUsed] = Source
      .actorRef[OutgoingMessage](BUFFER_SIZE, OverflowStrategy.backpressure)
      .mapMaterializedValue { socket =>
        clientEndpoint ! Attach(socket)
        NotUsed
      }
      .map {
        case OutgoingMessage(jsContent) => TextMessage(jsContent)
      }

    Flow.fromSinkAndSource(incomingMessages, outgoingMessages)
  }
}

object WebSocketHandler {
  private val BUFFER_SIZE: Int = 10
} 

//设置package包名称以及导入依赖的类
package net.ruippeixotog.scalafbp.ws

import akka.event.slf4j.SLF4JLogging
import akka.http.scaladsl.model.ws.{ BinaryMessage, Message, TextMessage }
import akka.stream.scaladsl.Flow

trait WsUtils extends SLF4JLogging {

  private[this] def printMessage(msg: Message): String = msg match {
    case TextMessage.Strict(text) => text
    case _: TextMessage.Streamed => "(streamed text message)"
    case _: BinaryMessage => "(binary message)"
  }

  def logWsMessages(id: String)(handler: Flow[Message, Message, Any]) = {
    Flow[Message]
      .map { msg => log.debug(s"[$id] IN : ${printMessage(msg)}"); msg }
      .via(handler)
      .map { msg => log.debug(s"[$id] OUT: ${printMessage(msg)}"); msg }
  }
} 

//设置package包名称以及导入依赖的类
package co.technius.chatty.server.service

import akka.actor.ActorSystem
import akka.stream.{ FlowShape, Materializer }
import akka.stream.scaladsl._
import akka.stream.actor.{ ActorPublisher, ActorSubscriber }
import akka.http.scaladsl.model.ws.{ BinaryMessage, Message, TextMessage }

class ChatService(implicit system: ActorSystem, mat: Materializer) {
  val roomActor = system.actorOf(RoomActor.props("default"), "defaultroom")
  val roomPub =
    Source
      .fromPublisher[String](ActorPublisher(roomActor))
      .map(msg => TextMessage(Source.single(msg)))

  def flow(name: String): Flow[Message, Message, Any] = {
    val userActor = system.actorOf(UserActor.props(name, roomActor))
    roomActor.tell(InternalProtocol.Join(name), userActor)
    Flow.fromGraph(GraphDSL.create() { implicit b =>
      import GraphDSL.Implicits._

      val userOut = b.add(Source.fromPublisher(ActorPublisher(userActor)))
      val userIn = b.add(Sink.fromSubscriber(ActorSubscriber(userActor)))

      val fromMessage = b.add(msgToStringFlow)
      val toMessage = b.add(Flow[String].map(msg => TextMessage(msg)))

      fromMessage ~> userIn
      userOut ~> toMessage

      FlowShape(fromMessage.in, toMessage.out)
    })
  }

  def msgToStringFlow: Flow[Message, String, Any] = Flow[Message].mapConcat {
    case TextMessage.Strict(msg) => msg :: Nil
    case tm: TextMessage =>
      tm.textStream.runWith(Sink.ignore)
      Nil
    case bm: BinaryMessage =>
      bm.dataStream.runWith(Sink.ignore)
      Nil
  }

} 

//设置package包名称以及导入依赖的类
package akkaviz.server

import akka.http.scaladsl.model.ws.{BinaryMessage, Message}
import akka.stream.scaladsl.Flow
import akka.util.ByteString
import akkaviz.protocol

trait ProtocolSerializationSupport {

  def protocolServerMessageToByteString: Flow[protocol.ApiServerMessage, ByteString, Any] = Flow[protocol.ApiServerMessage].map {
    msg => ByteString(protocol.IO.write(msg))
  }

  def websocketMessageToClientMessage: Flow[Message, protocol.ApiClientMessage, _] = Flow[Message].collect {
    case BinaryMessage.Strict(msg) =>
      protocol.IO.readClient(msg.asByteBuffer)
  }

}

object ProtocolSerializationSupport extends ProtocolSerializationSupport 

//设置package包名称以及导入依赖的类
package akkaviz.server

import akka.actor.ActorSystem
import akka.http.scaladsl.model.ws.BinaryMessage
import akka.stream.ActorMaterializer
import akka.stream.scaladsl.{Sink, Source}
import akka.util.ByteString
import akkaviz.protocol.{IO, Killed, SetEnabled}
import org.scalatest.concurrent.ScalaFutures
import org.scalatest.{FunSuite, Matchers}

import scala.concurrent.duration._

class ProtocolSerializationSupportTest extends FunSuite with ScalaFutures with Matchers {

  private[this] implicit val system = ActorSystem()
  private[this] implicit val materializer = ActorMaterializer()

  override implicit val patienceConfig: PatienceConfig = PatienceConfig(10.seconds)

  import ProtocolSerializationSupport._

  test("websocketMessageToClientMessage") {
    val msg = SetEnabled(true)
    val wsMessage = BinaryMessage(ByteString(IO.write(msg)))
    val res = Source.single(wsMessage).via(websocketMessageToClientMessage).runWith(Sink.head)
    whenReady(res) {
      _ shouldBe msg
    }
  }

  test("protocolServerMessageToByteString") {
    val msg = Killed("ref")
    val res = Source.single(msg).via(protocolServerMessageToByteString).runWith(Sink.head)
    whenReady(res) {
      serialized =>
        IO.readServer(serialized.asByteBuffer) shouldBe msg
    }
  }

}