package com.thatdot.quine.app.ingest

import java.nio.charset.Charset

import scala.compat.ExecutionContexts
import scala.concurrent.duration.DurationInt
import scala.util.{Failure, Success}

import akka.NotUsed
import akka.actor.ActorSystem
import akka.http.scaladsl.Http
import akka.http.scaladsl.model.ws.{
  BinaryMessage,
  InvalidUpgradeResponse,
  Message,
  TextMessage,
  ValidUpgrade,
  WebSocketRequest
}
import akka.http.scaladsl.settings.ClientConnectionSettings
import akka.stream.KillSwitches
import akka.stream.contrib.{SwitchMode, Valve}
import akka.stream.scaladsl.{Flow, Keep, Source}
import akka.util.ByteString

import com.typesafe.scalalogging.LazyLogging

import com.thatdot.quine.app.ControlSwitches
import com.thatdot.quine.app.ingest.WebsocketSimpleStartup.UpgradeFailedException
import com.thatdot.quine.app.ingest.serialization.ImportFormat
import com.thatdot.quine.app.routes.IngestMeter
import com.thatdot.quine.graph.CypherOpsGraph
import com.thatdot.quine.graph.MasterStream.{IngestSrcExecToken, IngestSrcType}
import com.thatdot.quine.routes.WebsocketSimpleStartupIngest
import com.thatdot.quine.routes.WebsocketSimpleStartupIngest.KeepaliveProtocol
object WebsocketSimpleStartup {
  class UpgradeFailedException(cause: Throwable)
      extends RuntimeException("Unable to upgrade to websocket connection", cause) {

    def this(cause: String) = this(new Throwable(cause))
  }
}
final case class WebsocketSimpleStartup(
  format: ImportFormat,
  wsUrl: String,
  initMessages: Seq[String],
  keepaliveProtocol: KeepaliveProtocol,
  parallelism: Int,
  charset: Charset,
  charsetTranscoder: Flow[ByteString, ByteString, NotUsed],
  meter: IngestMeter,
  initialSwitchMode: SwitchMode
) extends LazyLogging {
  @throws[IllegalArgumentException]("When the provided url is invalid")
  def stream(implicit graph: CypherOpsGraph): IngestSrcType[ControlSwitches] = {
    implicit val system: ActorSystem = graph.system
    val execToken = IngestSrcExecToken(s"WebSocket: $wsUrl")
    // NB Instead of killing this source with the downstream KillSwitch, we could switch this Source.never to a
    // Source.maybe, completing it with None to kill the connection -- this is closer to the docs for
    // webSocketClientFlow
    val outboundMessages = Source.fromIterator(() => initMessages.iterator).map(TextMessage(_)).concat(Source.never)
    val baseHttpClientSettings = ClientConnectionSettings(system)
    // Copy (and potentially tweak) baseHttpClientSettings for websockets usage
    val httpClientSettings = keepaliveProtocol match {
      case WebsocketSimpleStartupIngest.PingPongInterval(intervalMillis) =>
        baseHttpClientSettings.withWebsocketSettings(
          baseHttpClientSettings.websocketSettings.withPeriodicKeepAliveMaxIdle(intervalMillis.millis)
        )
      case WebsocketSimpleStartupIngest.SendMessageInterval(message, intervalMillis) =>
        baseHttpClientSettings.withWebsocketSettings(
          baseHttpClientSettings.websocketSettings
            .withPeriodicKeepAliveMaxIdle(intervalMillis.millis)
            .withPeriodicKeepAliveData(() => ByteString(message, charset))
        )
      case WebsocketSimpleStartupIngest.NoKeepalive => baseHttpClientSettings
    }
    val wsFlow = Http().webSocketClientFlow(
      WebSocketRequest(wsUrl),
      settings = httpClientSettings
    )

    val (websocketUpgraded, websocketSource) = outboundMessages
      .viaMat(wsFlow)(Keep.right)
      .preMaterialize()

    Source
      .futureSource(websocketUpgraded.transform {
        // if the websocket upgrade fails, return an already-failed Source
        case Success(InvalidUpgradeResponse(_, cause)) => Failure(new UpgradeFailedException(cause))
        case Failure(ex) => Failure(new UpgradeFailedException(ex))
        // the websocket upgrade succeeded: proceed with setting up the ingest stream source
        case Success(ValidUpgrade(_, _)) =>
          Success(
            websocketSource
              .viaMat(KillSwitches.single)(Keep.right)
              .viaMat(Valve[Message](initialSwitchMode))(Keep.both)
              .flatMapConcat { msg =>
                val msgBytesSrc: Source[ByteString, _] =
                  msg match {
                    case textMessage: TextMessage =>
                      textMessage.textStream
                        .fold("")(_ + _)
                        .map(ByteString.fromString(_, charset))
                    case m: BinaryMessage =>
                      m.dataStream.fold(ByteString.empty)(_ concat _)
                  }
                msgBytesSrc.mapConcat { bytes =>
                  meter.mark(bytes.length)
                  format
                    .importMessageSafeBytes(bytes.toArray, graph.isSingleHost)
                    .fold(
                      { err =>

                        logger.warn(
                          s"Failed to deserialize a Websockets message (logged at INFO level). Skipping this message."
                        )
                        logger.info(
                          s"Failed to deserialize Websockets message UTF-8 representation: ${bytes.utf8String}",
                          err
                        )
                        List.empty
                      },
                      List(_)
                    )
                }
              }
              .via(graph.ingestThrottleFlow)
              .mapAsyncUnordered(parallelism)(format.writeToGraph(graph, _))
              .map(_ => execToken)
              .watchTermination() { case ((killSwitch, valveSwitch), doneFut) =>
                valveSwitch.map(v => ControlSwitches(killSwitch, v, doneFut))(ExecutionContexts.parasitic)
              }
          )
      }(ExecutionContexts.parasitic))
      .mapMaterializedValue(_.flatten)

  }
}
