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

package org.apache.pekko.stream.connectors.ftp
package impl

import org.apache.pekko
import pekko.stream.impl.Stages.DefaultAttributes.IODispatcher
import pekko.stream.stage.{ GraphStageWithMaterializedValue, InHandler, OutHandler }
import pekko.stream.{
  Attributes,
  IOOperationIncompleteException,
  IOResult,
  Inlet,
  Outlet,
  Shape,
  SinkShape,
  SourceShape
}
import pekko.util.ByteString
import pekko.util.ByteString.ByteString1C

import scala.concurrent.{ Future, Promise }
import java.io.{ IOException, InputStream, OutputStream }
import pekko.annotation.InternalApi

import scala.util.control.NonFatal

/**
 * INTERNAL API
 */
@InternalApi
private[ftp] trait FtpIOGraphStage[FtpClient, S <: RemoteFileSettings, Sh <: Shape]
    extends GraphStageWithMaterializedValue[Sh, Future[IOResult]] {

  def name: String

  def path: String

  def connectionSettings: S

  implicit def ftpClient: S => FtpClient

  val ftpLike: FtpLike[FtpClient, S]

  override def initialAttributes: Attributes =
    super.initialAttributes and Attributes.name(name) and IODispatcher

  override def shape: Sh
}

/**
 * INTERNAL API
 */
@InternalApi
private[ftp] trait FtpIOSourceStage[FtpClient, S <: RemoteFileSettings]
    extends FtpIOGraphStage[FtpClient, S, SourceShape[ByteString]] {

  def chunkSize: Int

  def offset: Long = 0L

  val shape: SourceShape[ByteString] = SourceShape(Outlet[ByteString](s"$name.out"))
  val out: Outlet[ByteString] = shape.outlets.head.asInstanceOf[Outlet[ByteString]]

  def createLogicAndMaterializedValue(inheritedAttributes: Attributes) = {

    val matValuePromise = Promise[IOResult]()

    val logic = new FtpGraphStageLogic[ByteString, FtpClient, S](shape, ftpLike, connectionSettings, ftpClient) {

      private var isOpt: Option[InputStream] = None
      private var readBytesTotal: Long = 0L

      setHandler(
        out,
        new OutHandler {
          def onPull(): Unit =
            try {
              readChunk() match {
                case Some(bs) =>
                  push(out, bs)
                case None =>
                  complete(out)
              }
            } catch {
              case NonFatal(e) =>
                failed = true
                matFailure(e)
                failStage(e)
            }
        }) // end of handler

      override def postStop(): Unit =
        try {
          isOpt.foreach { os =>
            try {
              os.close()
              graphStageFtpLike match {
                case cfo: CommonFtpOperations =>
                  if (!cfo.completePendingCommand(handler.get.asInstanceOf[cfo.Handler]))
                    throw new IOException("File transfer failed.")
                case _ =>
              }
            } catch {
              case e: IOException =>
                matFailure(e)
                // If we failed, we have to expect the stream might already be dead
                // so swallow the IOException
                if (!failed) throw e
              case NonFatal(e) =>
                matFailure(e)
                throw e
            }
          }
        } finally {
          super.postStop()
        }

      protected[this] def doPreStart(): Unit =
        isOpt = graphStageFtpLike match {
          case ur: UnconfirmedReads =>
            withUnconfirmedReads(ur)
          case ro: RetrieveOffset =>
            Some(ro.retrieveFileInputStream(path, handler.get.asInstanceOf[ro.Handler], offset).get)
          case _ =>
            Some(graphStageFtpLike.retrieveFileInputStream(path, handler.get).get)
        }

      private def withUnconfirmedReads(
          ftpLikeWithUnconfirmedReads: FtpLike[FtpClient, S] with UnconfirmedReads): Option[InputStream] =
        connectionSettings match {
          case s: SftpSettings =>
            Some(
              ftpLikeWithUnconfirmedReads
                .retrieveFileInputStream(path,
                  handler.get.asInstanceOf[ftpLikeWithUnconfirmedReads.Handler],
                  offset,
                  s.maxUnconfirmedReads)
                .get)
          case _ => None
        }

      protected[this] def matSuccess(): Boolean =
        matValuePromise.trySuccess(IOResult.createSuccessful(readBytesTotal))

      protected[this] def matFailure(t: Throwable): Boolean =
        matValuePromise.tryFailure(new IOOperationIncompleteException(readBytesTotal, t))

      /** BLOCKING I/O READ */
      private def readChunk() = {
        def read(arr: Array[Byte]) =
          isOpt.flatMap { is =>
            val readBytes = is.read(arr)
            if (readBytes > -1) Some(readBytes)
            else None
          }
        val arr = Array.ofDim[Byte](chunkSize)
        read(arr).map { readBytes =>
          readBytesTotal += readBytes
          if (readBytes == chunkSize)
            ByteString1C(arr)
          else
            ByteString1C(arr).take(readBytes)
        }
      }

    } // end of stage logic

    (logic, matValuePromise.future)
  }

}

/**
 * INTERNAL API
 */
@InternalApi
private[ftp] trait FtpIOSinkStage[FtpClient, S <: RemoteFileSettings]
    extends FtpIOGraphStage[FtpClient, S, SinkShape[ByteString]] {

  def append: Boolean

  val shape: SinkShape[ByteString] = SinkShape(Inlet[ByteString](s"$name.in"))
  val in: Inlet[ByteString] = shape.inlets.head.asInstanceOf[Inlet[ByteString]]

  def createLogicAndMaterializedValue(inheritedAttributes: Attributes) = {

    val matValuePromise = Promise[IOResult]()

    val logic = new FtpGraphStageLogic[ByteString, FtpClient, S](shape, ftpLike, connectionSettings, ftpClient) {

      private var osOpt: Option[OutputStream] = None
      private var writtenBytesTotal: Long = 0L

      setHandler(
        in,
        new InHandler {
          override def onPush(): Unit =
            try {
              write(grab(in))
              pull(in)
            } catch {
              case NonFatal(e) =>
                failed = true
                matFailure(e)
                failStage(e)
            }

          override def onUpstreamFailure(exception: Throwable): Unit = {
            matFailure(exception)
            failed = true
            super.onUpstreamFailure(exception)
          }
        }) // end of handler

      override def postStop(): Unit =
        try {
          osOpt.foreach { os =>
            try {
              os.close()
              graphStageFtpLike match {
                case cfo: CommonFtpOperations =>
                  if (!cfo.completePendingCommand(handler.get.asInstanceOf[cfo.Handler]))
                    throw new IOException("File transfer failed.")
                case _ =>
              }
            } catch {
              case e: IOException =>
                matFailure(e)
                // If we failed, we have to expect the stream might already be dead
                // so swallow the IOException
                if (!failed) throw e
              case NonFatal(e) =>
                matFailure(e)
                throw e
            }
          }
        } finally {
          super.postStop()
        }

      protected[this] def doPreStart(): Unit = {
        osOpt = Some(graphStageFtpLike.storeFileOutputStream(path, handler.get, append).get)
        pull(in)
      }

      protected[this] def matSuccess(): Boolean =
        matValuePromise.trySuccess(IOResult.createSuccessful(writtenBytesTotal))

      protected[this] def matFailure(t: Throwable): Boolean =
        matValuePromise.tryFailure(new IOOperationIncompleteException(writtenBytesTotal, t))

      /** BLOCKING I/O WRITE */
      private def write(bytes: ByteString) =
        osOpt.foreach { os =>
          os.write(bytes.toArray)
          writtenBytesTotal += bytes.size
        }

    } // end of stage logic

    (logic, matValuePromise.future)
  }

}

/**
 * INTERNAL API
 */
@InternalApi
private[ftp] trait FtpMoveSink[FtpClient, S <: RemoteFileSettings]
    extends GraphStageWithMaterializedValue[SinkShape[FtpFile], Future[IOResult]] {
  val connectionSettings: S
  val ftpClient: S => FtpClient
  val ftpLike: FtpLike[FtpClient, S]
  val destinationPath: FtpFile => String
  val in: Inlet[FtpFile] = Inlet("FtpMvSink")

  def shape: SinkShape[FtpFile] = SinkShape(in)

  def createLogicAndMaterializedValue(inheritedAttributes: Attributes) = {
    val matValuePromise = Promise[IOResult]()
    var numberOfMovedFiles = 0

    val logic = new FtpGraphStageLogic[FtpFile, FtpClient, S](shape, ftpLike, connectionSettings, ftpClient) {
      {
        setHandler(
          in,
          new InHandler {
            override def onPush(): Unit = {
              try {
                val sourcePath = grab(in)
                graphStageFtpLike.move(sourcePath.path, destinationPath(sourcePath), handler.get)
                numberOfMovedFiles = numberOfMovedFiles + 1
                pull(in)
              } catch {
                case NonFatal(e) =>
                  failed = true
                  matFailure(e)
                  failStage(e)
              }
            }

            override def onUpstreamFailure(exception: Throwable): Unit = {
              matFailure(exception)
              failed = true
              super.onUpstreamFailure(exception)
            }
          })
      }

      protected[this] def doPreStart(): Unit = pull(in)

      protected[this] def matSuccess(): Boolean =
        matValuePromise.trySuccess(IOResult.createSuccessful(numberOfMovedFiles))

      protected[this] def matFailure(t: Throwable): Boolean =
        matValuePromise.tryFailure(t)
    } // end of stage logic

    (logic, matValuePromise.future)
  }
}

/**
 * INTERNAL API
 */
@InternalApi
private[ftp] trait FtpRemoveSink[FtpClient, S <: RemoteFileSettings]
    extends GraphStageWithMaterializedValue[SinkShape[FtpFile], Future[IOResult]] {
  val connectionSettings: S
  val ftpClient: S => FtpClient
  val ftpLike: FtpLike[FtpClient, S]
  val in: Inlet[FtpFile] = Inlet("FtpRmSink")

  def shape: SinkShape[FtpFile] = SinkShape(in)

  def createLogicAndMaterializedValue(inheritedAttributes: Attributes) = {
    val matValuePromise = Promise[IOResult]()
    var numberOfRemovedFiles = 0
    val logic = new FtpGraphStageLogic[Unit, FtpClient, S](shape, ftpLike, connectionSettings, ftpClient) {
      {
        setHandler(
          in,
          new InHandler {
            override def onPush(): Unit = {
              try {
                graphStageFtpLike.remove(grab(in).path, handler.get)
                numberOfRemovedFiles = numberOfRemovedFiles + 1
                pull(in)
              } catch {
                case NonFatal(e) =>
                  failed = true
                  matFailure(e)
                  failStage(e)
              }
            }

            override def onUpstreamFailure(exception: Throwable): Unit = {
              matFailure(exception)
              failed = true
              super.onUpstreamFailure(exception)
            }
          })
      }

      protected[this] def doPreStart(): Unit = pull(in)

      protected[this] def matSuccess(): Boolean =
        matValuePromise.trySuccess(IOResult.createSuccessful(numberOfRemovedFiles))

      protected[this] def matFailure(t: Throwable): Boolean =
        matValuePromise.tryFailure(t)
    } // end of stage logic

    (logic, matValuePromise.future)
  }
}
