/*
 * 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) 2020-2022 Lightbend Inc. <https://www.lightbend.com>
 */

package org.apache.pekko.stream.impl.fusing

import org.apache.pekko
import pekko.annotation.InternalApi
import pekko.dispatch.ExecutionContexts
import pekko.stream.Attributes.SourceLocation
import pekko.stream.impl.Stages.DefaultAttributes
import pekko.stream.{
  AbruptStageTerminationException,
  Attributes,
  FlowShape,
  Inlet,
  NeverMaterializedException,
  Outlet
}
import pekko.stream.scaladsl.{ Flow, Keep, Source }
import pekko.stream.stage.{ GraphStageLogic, GraphStageWithMaterializedValue, InHandler, OutHandler }
import pekko.util.OptionVal

import scala.concurrent.{ Future, Promise }
import scala.util.control.NonFatal
import scala.util.{ Failure, Success, Try }

@InternalApi private[pekko] final class FutureFlow[In, Out, M](futureFlow: Future[Flow[In, Out, M]])
    extends GraphStageWithMaterializedValue[FlowShape[In, Out], Future[M]] {

  val in = Inlet[In]("FutureFlow.in")
  val out = Outlet[Out]("FutureFlow.out")

  override protected def initialAttributes: Attributes =
    DefaultAttributes.futureFlow and SourceLocation.forLambda(futureFlow)

  override val shape: FlowShape[In, Out] = FlowShape(in, out)

  override def createLogicAndMaterializedValue(inheritedAttributes: Attributes): (GraphStageLogic, Future[M]) = {
    val propagateToNestedMaterialization =
      inheritedAttributes
        .mandatoryAttribute[Attributes.NestedMaterializationCancellationPolicy]
        .propagateToNestedMaterialization
    val innerMatValue = Promise[M]()
    val logic = new GraphStageLogic(shape) {

      // seems like we must set handlers BEFORE preStart
      setHandlers(in, out, Initializing)

      override def preStart(): Unit = {
        futureFlow.value match {
          case Some(tryFlow) =>
            Initializing.onFuture(tryFlow)
          case None =>
            val cb = getAsyncCallback(Initializing.onFuture)
            futureFlow.onComplete(cb.invoke)(ExecutionContexts.parasitic)
            // in case both ports are closed before future completion
            setKeepGoing(true)
        }
      }

      override def postStop(): Unit = {
        if (!innerMatValue.isCompleted)
          innerMatValue.failure(new AbruptStageTerminationException(this))
      }

      object Initializing extends InHandler with OutHandler {
        // we don't expect a push since we bever pull upstream during initialization
        override def onPush(): Unit = throw new IllegalStateException("unexpected push during initialization")

        var upstreamFailure = OptionVal.none[Throwable]

        override def onUpstreamFailure(ex: Throwable): Unit = {
          upstreamFailure = OptionVal.Some(ex)
        }

        // will later be propagated to the materialized flow (by examining isClosed(in))
        override def onUpstreamFinish(): Unit = {}

        // will later be propagated to the materialized flow (by examining isAvailable(out))
        override def onPull(): Unit = {}

        var downstreamCause = OptionVal.none[Throwable]

        override def onDownstreamFinish(cause: Throwable): Unit =
          if (propagateToNestedMaterialization) {
            downstreamCause = OptionVal.Some(cause)
          } else {
            innerMatValue.failure(new NeverMaterializedException(cause))
            cancelStage(cause)
          }

        def onFuture(futureRes: Try[Flow[In, Out, M]]) = futureRes match {
          case Failure(exception) =>
            setKeepGoing(false)
            innerMatValue.failure(new NeverMaterializedException(exception))
            failStage(exception)
          case Success(flow) =>
            // materialize flow, connect inlet and outlet, feed with potential events and set handlers
            connect(flow)
            setKeepGoing(false)
        }

        def connect(flow: Flow[In, Out, M]): Unit = {
          val subSource = new SubSourceOutlet[In]("FutureFlow.subIn")
          val subSink = new SubSinkInlet[Out]("FutureFlow.subOut")

          subSource.setHandler {
            new OutHandler {
              override def onPull(): Unit = if (!isClosed(in)) tryPull(in)
              override def onDownstreamFinish(cause: Throwable): Unit = if (!isClosed(in)) cancel(in, cause)
            }
          }
          subSink.setHandler {
            new InHandler {
              override def onPush(): Unit = push(out, subSink.grab())
              override def onUpstreamFinish(): Unit = complete(out)
              override def onUpstreamFailure(ex: Throwable): Unit = fail(out, ex)
            }
          }
          try {
            val matVal = subFusingMaterializer.materialize(
              Source.fromGraph(subSource.source).viaMat(flow)(Keep.right).to(subSink.sink),
              inheritedAttributes)
            innerMatValue.success(matVal)
            upstreamFailure match {
              case OptionVal.Some(ex) => subSource.fail(ex)
              case _                  => if (isClosed(in)) subSource.complete()
            }
            downstreamCause match {
              case OptionVal.Some(cause) => subSink.cancel(cause)
              case _                     => if (isAvailable(out)) subSink.pull()
            }
            setHandlers(in, out,
              new InHandler with OutHandler {
                override def onPull(): Unit = subSink.pull()
                override def onDownstreamFinish(cause: Throwable): Unit = subSink.cancel(cause)
                override def onPush(): Unit = subSource.push(grab(in))
                override def onUpstreamFinish(): Unit = subSource.complete()
                override def onUpstreamFailure(ex: Throwable): Unit = subSource.fail(ex)
              })
          } catch {
            case NonFatal(ex) =>
              innerMatValue.failure(new NeverMaterializedException(ex))
              failStage(ex)
          }
        }
      }
    }
    (logic, innerMatValue.future)
  }
}
