WKRemoteDataStoreController.scala

package controllers

import com.scalableminds.util.accesscontext.{AuthorizedAccessContext, DBAccessContext, GlobalAccessContext}
import com.scalableminds.util.objectid.ObjectId
import com.scalableminds.util.time.Instant
import com.scalableminds.util.tools.Fox
import com.scalableminds.webknossos.datastore.controllers.JobExportProperties
import com.scalableminds.webknossos.datastore.helpers.UPath
import com.scalableminds.webknossos.datastore.models.UnfinishedUpload
import com.scalableminds.webknossos.datastore.models.datasource.{
  DataSource,
  DataSourceId,
  DataSourceStatus,
  LayerAttachmentType,
  UnusableDataSource
}
import com.scalableminds.webknossos.datastore.services.{DataSourcePathInfo, DataStoreStatus}
import com.scalableminds.webknossos.datastore.services.uploading.{
  AttachmentUploadAdditionalInfo,
  AttachmentUploadInfo,
  DatasetUploadAdditionalInfo,
  DatasetUploadInfo,
  MagUploadAdditionalInfo,
  MagUploadInfo,
  ReportAttachmentUploadParameters,
  ReportDatasetUploadParameters,
  ReportMagUploadParameters
}
import com.typesafe.scalalogging.LazyLogging
import models.dataset._
import models.dataset.credential.CredentialDAO
import models.job.{JobDAO, JobService}
import models.organization.{OrganizationDAO, OrganizationService}
import models.storage.UsedStorageService
import models.team.TeamDAO
import models.user.UserDAO
import play.api.i18n.Messages
import play.api.libs.json.Json
import play.api.mvc.{Action, AnyContent, PlayBodyParsers}
import security.{WebknossosBearerTokenAuthenticatorService, WkSilhouetteEnvironment}

import scala.concurrent.duration.DurationInt
import javax.inject.Inject
import scala.concurrent.ExecutionContext

class WKRemoteDataStoreController @Inject()(
    datasetService: DatasetService,
    dataStoreService: DataStoreService,
    dataStoreDAO: DataStoreDAO,
    organizationDAO: OrganizationDAO,
    usedStorageService: UsedStorageService,
    organizationService: OrganizationService,
    layerToLinkService: LayerToLinkService,
    datasetDAO: DatasetDAO,
    userDAO: UserDAO,
    teamDAO: TeamDAO,
    jobDAO: JobDAO,
    datasetMagDAO: DatasetMagDAO,
    datasetAttachmentDAO: DatasetLayerAttachmentDAO,
    uploadToPathsService: UploadToPathsService,
    jobService: JobService,
    credentialDAO: CredentialDAO,
    wkSilhouetteEnvironment: WkSilhouetteEnvironment)(implicit ec: ExecutionContext, bodyParsers: PlayBodyParsers)
    extends Controller
    with LazyLogging {

  val bearerTokenService: WebknossosBearerTokenAuthenticatorService =
    wkSilhouetteEnvironment.combinedAuthenticatorService.tokenAuthenticatorService

  def reserveDatasetUpload(name: String, key: String, token: String): Action[DatasetUploadInfo] =
    Action.async(validateJson[DatasetUploadInfo]) { implicit request =>
      dataStoreService.validateAccess(name, key) { dataStore =>
        val uploadInfo = request.body
        for {
          user <- bearerTokenService.userForToken(token) ~> FORBIDDEN
          organization <- organizationDAO.findOne(uploadInfo.organizationId)(GlobalAccessContext) ?~> Messages(
            "organization.notFound",
            uploadInfo.organizationId) ~> NOT_FOUND
          _ <- organizationService.assertUsedStorageNotExceeded(
            organization,
            uploadInfo.resumableUploadInfo.totalFileSizeInBytes) ?~> "dataset.upload.storageExceeded" ~> FORBIDDEN
          _ <- Fox.fromBool(organization._id == user._organization) ?~> "notAllowed" ~> FORBIDDEN
          _ <- datasetService.assertValidDatasetName(uploadInfo.datasetName)
          _ <- Fox.fromBool(dataStore.onlyAllowedOrganization.forall(_ == organization._id)) ?~> "dataset.upload.Datastore.restricted"
          _ <- Fox.serialCombined(uploadInfo.layersToLink.getOrElse(List.empty))(l =>
            layerToLinkService.validateLayerToLink(l, user)) ?~> "dataset.upload.invalidLinkedLayers"
          _ <- Fox.runIf(request.body.requireUniqueName.getOrElse(false))(
            datasetService.assertNewDatasetNameUnique(request.body.datasetName, organization._id))
          preliminaryDataSource = UnusableDataSource(DataSourceId("", ""), None, DataSourceStatus.notYetUploaded)
          dataset <- datasetService.createAndSetUpDataset(
            uploadInfo.datasetName,
            dataStore,
            preliminaryDataSource,
            uploadInfo.folderId,
            user,
            isVirtual = uploadInfo.isVirtual.getOrElse(true),
            creationType = DatasetCreationType.Upload
          ) ?~> "dataset.upload.creation.failed"
          _ <- datasetService.addInitialTeams(dataset, uploadInfo.initialTeamIds, user)(AuthorizedAccessContext(user))
          additionalInfo = DatasetUploadAdditionalInfo(dataset._id, dataset.directoryName)
        } yield Ok(Json.toJson(additionalInfo))
      }
    }

  def reserveMagUpload(name: String, key: String, token: String): Action[MagUploadInfo] =
    Action.async(validateJson[MagUploadInfo]) { implicit request =>
      dataStoreService.validateAccess(name, key) { dataStore =>
        // DS write access was asserted already at this point.
        for {
          user <- bearerTokenService.userForToken(token)
          dataset <- datasetDAO.findOne(request.body.datasetId)(AuthorizedAccessContext(user))
          _ <- Fox.fromBool(dataset.isVirtual) ?~> "dataset.reserveMagUpload.notVirtual"
          (dataSource, dataLayer) <- datasetService.getDataSourceAndLayerFor(dataset, request.body.layerName)
          _ <- Fox.fromBool(!dataLayer.mags.exists(_.mag.maxDim == request.body.mag.mag.maxDim)) ?~> s"New mag ${request.body.mag.mag} conflicts with existing mag of the layer."
          _ <- Fox.fromBool(dataset._dataStore == dataStore.name) ?~> "Cannot upload mag to existing dataset via different datastore."
          _ <- uploadToPathsService.handleExistingPendingMag(dataset,
                                                             request.body.layerName,
                                                             request.body.mag.mag,
                                                             request.body.overwritePending)
          _ <- datasetMagDAO.insertWithUploadPending(request.body.datasetId,
                                                     request.body.layerName,
                                                     request.body.mag.mag,
                                                     request.body.mag.axisOrder,
                                                     request.body.mag.channelIndex)
        } yield Ok(Json.toJson(MagUploadAdditionalInfo(dataSource.id)))
      }
    }

  def reserveAttachmentUpload(name: String, key: String, token: String): Action[AttachmentUploadInfo] =
    Action.async(validateJson[AttachmentUploadInfo]) { implicit request =>
      dataStoreService.validateAccess(name, key) { dataStore =>
        // DS write access was asserted already at this point.
        for {
          user <- bearerTokenService.userForToken(token)
          dataset <- datasetDAO.findOne(request.body.datasetId)(AuthorizedAccessContext(user))
          _ <- Fox.fromBool(dataset.isVirtual) ?~> "dataset.reserveAttachmentUpload.notVirtual"
          (dataSource, dataLayer) <- datasetService.getDataSourceAndLayerFor(dataset, request.body.layerName)
          isSingletonAttachment = LayerAttachmentType.isSingletonAttachment(request.body.attachmentType)
          existsError = if (isSingletonAttachment) "attachment.singleton.alreadyFilled" else "attachment.name.taken"
          existingAttachmentOpt = dataLayer.attachments.flatMap(
            _.getByTypeAndNameAlwaysReturnSingletons(request.body.attachmentType, request.body.attachment.name))
          _ <- Fox.fromBool(existingAttachmentOpt.isEmpty) ?~> existsError
          _ <- Fox.fromBool(dataset._dataStore == dataStore.name) ?~> "Cannot upload attachment to existing dataset via different datastore."
          dummyAttachmentPath <- UPath.fromString("<pending upload>").toFox
          _ <- uploadToPathsService.handleExistingPendingAttachment(dataset,
                                                                    request.body.layerName,
                                                                    request.body.attachmentType,
                                                                    request.body.attachment.name,
                                                                    request.body.overwritePending)
          _ <- datasetAttachmentDAO.insertWithUploadPending(
            request.body.datasetId,
            request.body.layerName,
            request.body.attachment.name,
            request.body.attachmentType,
            request.body.attachment.dataFormat,
            dummyAttachmentPath
          )
        } yield Ok(Json.toJson(AttachmentUploadAdditionalInfo(dataSource.id)))
      }
    }

  def getUnfinishedDatasetUploadsForUser(name: String,
                                         key: String,
                                         token: String,
                                         organizationId: String): Action[AnyContent] =
    Action.async { implicit request =>
      dataStoreService.validateAccess(name, key) { _ =>
        for {
          user <- bearerTokenService.userForToken(token) ~> FORBIDDEN
          organization <- organizationDAO.findOne(organizationId)(GlobalAccessContext) ?~> Messages(
            "organization.notFound",
            user._organization) ~> NOT_FOUND
          _ <- Fox.fromBool(organization._id == user._organization) ?~> "notAllowed" ~> FORBIDDEN
          datasets <- datasetService.getAllUnfinishedDatasetUploadsOfUser(user._id, user._organization)(
            GlobalAccessContext) ?~> "dataset.upload.couldNotLoadUnfinishedUploads"
          teamIdsPerDataset <- Fox.combined(datasets.map(dataset => teamDAO.findAllowedTeamIdsForDataset(dataset.id)))
          unfinishedUploads = datasets.zip(teamIdsPerDataset).map {
            case (d, teamIds) =>
              UnfinishedUpload("<filled-in by datastore>",
                               d.dataSourceId,
                               d.name,
                               d.folderId.toString,
                               d.created,
                               None, // Filled by datastore.
                               teamIds.map(_.toString))
          }
        } yield Ok(Json.toJson(unfinishedUploads))
      }
    }

  def reportDatasetUpload(name: String,
                          key: String,
                          token: String,
                          datasetId: ObjectId): Action[ReportDatasetUploadParameters] =
    Action.async(validateJson[ReportDatasetUploadParameters]) { implicit request =>
      implicit val ctx: DBAccessContext = GlobalAccessContext
      dataStoreService.validateAccess(name, key) { _ =>
        for {
          user <- bearerTokenService.userForToken(token) ~> FORBIDDEN
          dataset <- datasetDAO.findOne(datasetId) ?~> Messages("dataset.notFound", datasetId) ~> NOT_FOUND
          _ = datasetService.trackNewDataset(dataset,
                                             user,
                                             request.body.needsConversion,
                                             request.body.datasetSizeBytes,
                                             addVariantLabel = "upload without conversion")
          dataSourceWithLinkedLayersOpt <- Fox.runOptional(request.body.dataSourceOpt) {
            implicit val ctx: DBAccessContext = AuthorizedAccessContext(user)
            layerToLinkService.addLayersToLinkToDataSource(_, request.body.layersToLink)
          }
          _ <- Fox.runOptional(dataSourceWithLinkedLayersOpt) { dataSource =>
            logger.info(s"Updating dataset $datasetId in database after upload reported from datastore $name.")
            datasetDAO.updateDataSource(datasetId,
                                        dataset._dataStore,
                                        dataSource.hashCode(),
                                        dataSource,
                                        isUsable = true)
          }
          updated <- datasetDAO.findOne(datasetId) ?~> Messages("dataset.notFound", datasetId) ~> NOT_FOUND
          _ <- Fox.runIf(!request.body.needsConversion)(usedStorageService.refreshStorageReportForDataset(updated))
          _ <- Fox.runIf(!request.body.needsConversion)(datasetService.scanRealpathsIfVirtual(updated))
          _ <- Fox.runIf(request.body.needsConversion) {
            for {
              voxelSize <- request.body.voxelSize.toFox ?~> "dataset.upload.needsConversion.missingVoxelSize"
              _ <- jobService.submitConvertToWkwJob(dataset, user, voxelSize)
            } yield ()
          }
        } yield Ok
      }
    }

  def reportMagUpload(name: String, key: String): Action[ReportMagUploadParameters] =
    Action.async(validateJson[ReportMagUploadParameters]) { implicit request =>
      dataStoreService.validateAccess(name, key) { _ =>
        for {
          dataset <- datasetDAO.findOne(request.body.datasetId)(GlobalAccessContext) ?~> Messages(
            "dataset.notFound",
            request.body.datasetId) ~> NOT_FOUND
          _ <- datasetMagDAO.findOneWithPendingUpload(request.body.datasetId,
                                                      request.body.layerName,
                                                      request.body.mag.mag) ?~> "dataset.finishMagUpload.notPending"
          _ <- request.body.mag.path.toFox ?~> "dataset.finishMagUpload.pathNotSet"
          _ <- datasetMagDAO.finishUpload(request.body.datasetId, request.body.layerName, request.body.mag)
          dataStoreClient <- datasetService.clientFor(dataset)(GlobalAccessContext)
          _ <- dataStoreClient.invalidateDatasetInDSCache(dataset._id)
          _ <- usedStorageService.refreshStorageReportForDataset(dataset)
        } yield Ok
      }
    }

  def reportAttachmentUpload(name: String, key: String): Action[ReportAttachmentUploadParameters] =
    Action.async(validateJson[ReportAttachmentUploadParameters]) { implicit request =>
      dataStoreService.validateAccess(name, key) { _ =>
        for {
          dataset <- datasetDAO.findOne(request.body.datasetId)(GlobalAccessContext) ?~> Messages(
            "dataset.notFound",
            request.body.datasetId) ~> NOT_FOUND
          _ <- datasetAttachmentDAO.findOneWithPendingUpload(
            request.body.datasetId,
            request.body.layerName,
            request.body.attachmentType,
            request.body.attachment.name) ?~> "dataset.finishAttachmentUpload.notPending"
          _ <- datasetAttachmentDAO.finishUpload(request.body.datasetId,
                                                 request.body.layerName,
                                                 request.body.attachmentType,
                                                 request.body.attachment)
          dataStoreClient <- datasetService.clientFor(dataset)(GlobalAccessContext)
          _ <- dataStoreClient.invalidateDatasetInDSCache(dataset._id)
          _ <- usedStorageService.refreshStorageReportForDataset(dataset)
        } yield Ok
      }
    }

  def statusUpdate(name: String, key: String): Action[DataStoreStatus] = Action.async(validateJson[DataStoreStatus]) {
    implicit request =>
      dataStoreService.validateAccess(name, key) { _ =>
        val okLabel = if (request.body.ok) "ok" else "not ok"
        logger.info(s"Status update from data store ‘$name’. Status $okLabel")
        for {
          _ <- dataStoreDAO.updateUrlByName(name, request.body.url)
        } yield Ok
      }
  }

  def updateAll(name: String, key: String, organizationId: Option[String]): Action[List[DataSource]] =
    Action.async(validateJson[List[DataSource]]) { implicit request =>
      dataStoreService.validateAccess(name, key) { dataStore =>
        implicit val ctx: DBAccessContext = GlobalAccessContext
        val dataSources = request.body
        for {
          before <- Instant.nowFox
          selectedOrgaLabel = organizationId.map(id => s"for organization $id").getOrElse("for all organizations")
          _ = logger.info(
            s"Received dataset list from datastore ${dataStore.name} $selectedOrgaLabel: " +
              s"${dataSources.count(_.isUsable)} active, ${dataSources.count(!_.isUsable)} inactive")
          existingIds <- datasetService.updateDataSources(dataStore, dataSources)
          _ <- datasetService.deactivateUnreportedDataSources(existingIds, dataStore, organizationId)
          _ = if (Instant.since(before) > (30 seconds))
            Instant.logSince(before,
                             s"Updating datasources from datastore ${dataStore.name} $selectedOrgaLabel",
                             logger)
        } yield Ok
      }
    }

  def updateOne(name: String, key: String): Action[DataSource] =
    Action.async(validateJson[DataSource]) { implicit request =>
      dataStoreService.validateAccess(name, key) { dataStore =>
        implicit val ctx: DBAccessContext = GlobalAccessContext
        for {
          _ <- datasetService.updateDataSources(dataStore, List(request.body))
        } yield Ok
      }
    }

  def updateRealPaths(name: String, key: String): Action[List[DataSourcePathInfo]] =
    Action.async(validateJson[List[DataSourcePathInfo]]) { implicit request =>
      dataStoreService.validateAccess(name, key) { _ =>
        for {
          _ <- datasetService.updateRealPaths(request.body)(GlobalAccessContext)
        } yield Ok
      }
    }

  /**
    * Called by the datastore after a dataset has been deleted on disk.
    */
  def deleteDataset(name: String, key: String): Action[ObjectId] = Action.async(validateJson[ObjectId]) {
    implicit request =>
      dataStoreService.validateAccess(name, key) { _ =>
        for {
          _ <- datasetService.deleteDatasetFromDB(request.body)
        } yield Ok
      }
  }

  def findDatasetId(name: String,
                    key: String,
                    datasetDirectoryName: String,
                    organizationId: String): Action[AnyContent] =
    Action.async { implicit request =>
      dataStoreService.validateAccess(name, key) { _ =>
        for {
          organization <- organizationDAO.findOne(organizationId)(GlobalAccessContext) ?~> Messages(
            "organization.notFound",
            organizationId) ~> NOT_FOUND
          dataset <- datasetDAO.findOneByNameAndOrganization(datasetDirectoryName, organization._id)(
            GlobalAccessContext) ?~> Messages("dataset.notFound", datasetDirectoryName)
        } yield Ok(Json.toJson(dataset._id))
      }
    }

  def getDataSource(name: String, key: String, datasetId: ObjectId): Action[AnyContent] =
    Action.async { implicit request =>
      dataStoreService.validateAccess(name, key) { _ =>
        for {
          dataset <- datasetDAO.findOne(datasetId)(GlobalAccessContext) ?~> Messages("dataset.notFound", datasetId) ~> NOT_FOUND
          dataSource <- datasetService.dataSourceFor(dataset)
        } yield Ok(Json.toJson(dataSource))
      }
    }

  def updateDataSource(name: String, key: String, datasetId: ObjectId): Action[DataSource] =
    Action.async(validateJson[DataSource]) { implicit request =>
      dataStoreService.validateAccess(name, key) { _ =>
        for {
          dataset <- datasetDAO.findOne(datasetId)(GlobalAccessContext) ?~> Messages("dataset.notFound", datasetId) ~> NOT_FOUND
          _ <- Fox.runIf(!dataset.isVirtual)(
            datasetDAO.updateDataSource(datasetId,
                                        name,
                                        request.body.hashCode(),
                                        request.body,
                                        isUsable = request.body.toUsable.isDefined)(GlobalAccessContext))
        } yield Ok
      }
    }

  def jobExportProperties(name: String, key: String, jobId: ObjectId): Action[AnyContent] = Action.async {
    implicit request =>
      dataStoreService.validateAccess(name, key) { _ =>
        for {
          job <- jobDAO.findOne(jobId)(GlobalAccessContext)
          jobOwner <- userDAO.findOne(job._owner)(GlobalAccessContext)
          organization <- organizationDAO.findOne(jobOwner._organization)(GlobalAccessContext)
          latestRunId <- job.latestRunId.toFox ?~> "job.notRun"
          exportFileName <- job.exportFileName.toFox ?~> "job.noExportFileName"
          jobExportProperties = JobExportProperties(jobId.toString, latestRunId, organization._id, exportFileName)
        } yield Ok(Json.toJson(jobExportProperties))
      }
  }

  def findCredential(name: String, key: String, credentialId: ObjectId): Action[AnyContent] = Action.async {
    implicit request =>
      dataStoreService.validateAccess(name, key) { _ =>
        for {
          credential <- credentialDAO.findOne(credentialId)
        } yield Ok(Json.toJson(credential))
      }
  }

}