import.go

// Licensed to NASA JPL under one or more contributor
// license agreements. See the NOTICE file distributed with
// this work for additional information regarding copyright
// ownership. NASA JPL licenses this file to you under
// the Apache License, Version 2.0 (the "License"); you may
// not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

// Exposes the interface of the dataset importer aka converter and selecting one automatically based on what
// files are in the folder being imported. The converter supports various formats as delivered by GDS or test
// instruments and this is inteded to be extendable further to other lab instruments and devices in future.
package dataimport

import (
	"fmt"
	"os"
	"path"
	"path/filepath"
	"strconv"
	"time"

	"github.com/pixlise/core/v4/api/dataimport/internal/converterSelector"
	"github.com/pixlise/core/v4/api/dataimport/internal/datasetArchive"
	"github.com/pixlise/core/v4/api/dataimport/internal/output"
	"github.com/pixlise/core/v4/api/filepaths"
	"github.com/pixlise/core/v4/core/fileaccess"
	"github.com/pixlise/core/v4/core/logger"
	"github.com/pixlise/core/v4/core/scan"
	protos "github.com/pixlise/core/v4/generated-protos"
	diffractionDetector "github.com/pixlise/diffraction-peak-detection/v2/detection"
	"go.mongodb.org/mongo-driver/mongo"
	"google.golang.org/protobuf/proto"
)

// All dataset conversions are started through here. This can contain multiple implementations
// for different scenarios, but internally it all runs the same way

// ImportFromArchive - Importing from dataset archive area. Calls ImportFromLocalFileSystem
// Returns:
// WorkingDir
// Saved dataset summary structure
// What changed (as a string), so caller can know what kind of notification to send (if any)
// IsUpdate flag
// Error (if any)
func ImportDataset(
	localFS fileaccess.FileAccess,
	remoteFS fileaccess.FileAccess,
	configBucket string,
	manualUploadBucket string,
	datasetBucket string,
	db *mongo.Database,
	datasetID string,
	log logger.ILogger,
	justArchived bool, // Set to true if a file was just saved to the archive prior to calling this. Affects notifications sent out
) (string, *protos.ScanItem, string, bool, error) {

	savedSummary := &protos.ScanItem{}

	workingDir, err := os.MkdirTemp("", "archive")
	if err != nil {
		return workingDir, savedSummary, "", false, err
	}

	// Read previously saved dataset summary file, so we have something to compare against to see what changes
	// we will need to notify on
	oldSummary, errOldSummary := scan.ReadScanItem(datasetID, db)
	if err != nil {
		// NOTE: we don't die here, we may be importing for the first time! Just log and continue
		//return workingDir, savedSummary, "", false, err
		log.Infof("Failed to import previous dataset summary file - assuming we're a new dataset...")
	}

	// Firstly, we download from the archive
	archive := datasetArchive.NewDatasetArchiveDownloader(remoteFS, localFS, log, datasetBucket, manualUploadBucket)
	localDownloadPath, localUnzippedPath, zipCount, err := archive.DownloadFromDatasetArchive(datasetID, workingDir)
	if err != nil {
		return workingDir, savedSummary, "", false, err
	}

	// If no zip files were loaded, maybe this dataset is a manually uploaded one, try to import from there instead
	if zipCount == 0 {
		log.Infof("No zip files found in archive, dataset may have been manually uploaded. Trying to download...")
		localDownloadPath, localUnzippedPath, err = archive.DownloadFromDatasetUploads(datasetID, workingDir)
		if err != nil {
			return workingDir, savedSummary, "", false, err
		}
	}

	// No obvious place to make this change right now, but pseudo-intensities have changed in flight software
	// and this is likely to go live in late 2023.
	pseudoVersion := ""

	iDatasetId, err := strconv.Atoi(datasetID)
	if err == nil {
		// Check if it requires the new pseudo-intensity ranges file
		if iDatasetId >= 297796101 {
			log.Infof("Detected need for new pseudo-intensity labels list, as this dataset is likely generated by updated FSW...")
			pseudoVersion = "-2023"
		}
	}

	localRangesPath, err := archive.DownloadPseudoIntensityRangesFile(configBucket, localDownloadPath, pseudoVersion)
	if err != nil {
		return workingDir, savedSummary, "", false, err
	}

	log.Infof("Downloading user customisation files...")

	err = archive.DownloadUserCustomisationsForDataset(datasetID, localUnzippedPath)
	if err != nil {
		return workingDir, savedSummary, "", false, err
	}

	// Now that we have data down, we can run the importer from local file system
	_, err = ImportFromLocalFileSystem(
		localFS,
		remoteFS,
		db,
		workingDir,
		localUnzippedPath,
		localRangesPath,
		datasetBucket,
		datasetID,
		log,
	)
	if err != nil {
		return workingDir, savedSummary, "", false, err
	}

	// Decide what notifications (if any) to send
	updatenotificationtype := "unknown"

	if errOldSummary == nil { // don't do this if the old summary couldn't be read!
		savedSummary, err = scan.ReadScanItem(datasetID, db)
		if err != nil {
			return workingDir, savedSummary, "", false, err
		}

		updatenotificationtype, err = getUpdateType(savedSummary, oldSummary)
		if err != nil {
			return workingDir, savedSummary, "", false, err
		}
	}

	return workingDir, savedSummary, updatenotificationtype, !justArchived && zipCount > 1, err
}

// ImportFromLocalFileSystem - As the name says, imports from directory on local file system
// Returns:
// Dataset ID (in case it was modified during conversion)
// Error (if there was one)
func ImportFromLocalFileSystem(
	localFS fileaccess.FileAccess,
	remoteFS fileaccess.FileAccess, // For uploading result
	db *mongo.Database,
	workingDir string, // Working dir, under which we may form our output dir
	localImportPath string, // Path on local file system with directory ready to import
	localPseudoIntensityRangesPath string, // Path on local file system
	datasetBucket string, // Where we import to
	datasetID string, // Dataset ID being imported. Some importers may need this, others (who have dataset ID in file names being imported) can verify it matches this expected one
	log logger.ILogger) (string, error) {

	// Pick an importer by inspecting the directory we're about to import from
	importer, err := converterSelector.SelectDataConverter(localFS, remoteFS, datasetBucket, localImportPath, log)

	if err != nil {
		return "", err
	}

	// Create an output directory
	outputScanPath, err := fileaccess.MakeEmptyLocalDirectory(workingDir, "output-"+filepaths.DatasetScansRoot)
	outputImagesPath, err := fileaccess.MakeEmptyLocalDirectory(workingDir, "output-"+filepaths.DatasetImagesRoot)

	if err != nil {
		return "", err
	}

	log.Infof("Running dataset converter...")
	data, contextImageSrcPath, err := importer.Import(localImportPath, localPseudoIntensityRangesPath, datasetID, log)
	if err != nil {
		return "", fmt.Errorf("Import failed: %v", err)
	}

	// Apply any overrides we may have
	customMetaFields, err := readLocalCustomMeta(log, localImportPath)
	if err != nil {
		return "", err
	}

	if len(customMetaFields.Title) > 0 && customMetaFields.Title != " " {
		log.Infof("Applying custom title: %v", customMetaFields.Title)
		data.Meta.Title = customMetaFields.Title
	}

	if len(customMetaFields.DefaultContextImage) > 0 {
		log.Infof("Applying custom default context image: %v", customMetaFields.DefaultContextImage)
		data.DefaultContextImage = customMetaFields.DefaultContextImage
	}

	// Form the output path
	outPath := filepath.Join(outputScanPath, data.DatasetID)

	log.Infof("Writing dataset file...")
	saver := output.PIXLISEDataSaver{}
	err = saver.Save(*data, contextImageSrcPath, outPath, filepath.Join(outputImagesPath, data.DatasetID), db, time.Now().Unix(), log)
	if err != nil {
		return "", fmt.Errorf("Failed to write dataset file: %v. Error: %v", outPath, err)
	}

	log.Infof("Running diffraction DB generator...")
	err = createPeakDiffractionDB(filepath.Join(outPath, filepaths.DatasetFileName), filepath.Join(outPath, filepaths.DiffractionDBFileName), log)

	if err != nil {
		return "", fmt.Errorf("Failed to run diffraction DB generator. Error: %v", err)
	}

	// Finally, copy scan files to scans, and images to images
	log.Infof("Copying generated dataset to bucket: %v...", datasetBucket)
	err = copyToBucket(remoteFS, data.DatasetID, outputScanPath, datasetBucket, filepaths.DatasetScansRoot, log)
	if err != nil {
		return "", fmt.Errorf("Error when copying dataset to bucket: %v. Error: %v", datasetBucket, err)
	}

	log.Infof("Copying images to bucket: %v...", datasetBucket)
	err = copyToBucket(remoteFS, data.DatasetID, outputImagesPath, datasetBucket, filepaths.DatasetImagesRoot, log)
	if err != nil {
		return "", fmt.Errorf("Error when copying dataset to bucket: %v. Error: %v", datasetBucket, err)
	}

	return data.DatasetID, nil
}

// createPeakDiffractoinDB - Use the diffraction engine to calculate the diffraction peaks
func createPeakDiffractionDB(datasetPath string, savepath string, jobLog logger.ILogger) error {
	localFS := fileaccess.FSAccess{}
	fileBytes, err := localFS.ReadObject("", datasetPath)
	if err != nil {
		jobLog.Errorf("Failed to open dataset \"%v\": \"%v\"", datasetPath, err)
		return err
	}

	// Now decode the data & return it
	protoParsed := &protos.Experiment{}
	err = proto.Unmarshal(fileBytes, protoParsed)
	if err != nil {
		jobLog.Errorf("Failed to read dataset \"%v\": \"%v\"", datasetPath, err)
		return err
	}

	jobLog.Infof("  Opened %v, got RTT: %v, title: \"%v\". Scanning for diffraction peaks...", datasetPath, protoParsed.Rtt, protoParsed.Title)

	datasetPeaks, err := diffractionDetector.ScanDataset(protoParsed)
	if err != nil {
		jobLog.Errorf("Error Encoundered During Scanning: %v", err)
		return err
	}

	jobLog.Infof("  Completed scan successfully")

	if savepath != "" {
		jobLog.Infof("  Saving diffraction db file: %v", savepath)
		diffractionPB := diffractionDetector.BuildDiffractionProtobuf(protoParsed, datasetPeaks)
		err := diffractionDetector.SaveDiffractionProtobuf(diffractionPB, savepath)
		if err != nil {
			jobLog.Errorf("Error Encoundered During Saving: %v", err)
			return err
		}

		jobLog.Infof("  Diffraction db saved successfully")
	}

	return nil
}

// Copies files to bucket
// NOTE: Assumes flat list of files, no folder structure!
func copyToBucket(remoteFS fileaccess.FileAccess, datasetID string, sourcePath string, destBucket string, destPath string, log logger.ILogger) error {
	var uploadError error

	err := filepath.Walk(sourcePath, func(sourcePath string, info os.FileInfo, err error) error {
		if !info.IsDir() {
			data, err := os.ReadFile(sourcePath)
			if err != nil {
				log.Errorf("Failed to read file for upload: %v", sourcePath)
				uploadError = err
			} else {
				sourceFile := filepath.Base(sourcePath)
				uploadPath := path.Join(destPath, datasetID, sourceFile)

				log.Infof("-Uploading: %v", sourcePath)
				log.Infof("---->to s3://%v/%v", destBucket, uploadPath)
				err = remoteFS.WriteObject(destBucket, uploadPath, data)

				if err != nil {
					log.Errorf("Failed to upload to s3://%v/%v: %v", destBucket, uploadPath, err)
					uploadError = err
				}
			}
		}
		return nil
	})

	if err != nil {
		return err
	}

	return uploadError
}

func getUpdateType(newSummary *protos.ScanItem, oldSummary *protos.ScanItem) (string, error) {
	for _, item := range newSummary.DataTypes {
		for _, oldItem := range oldSummary.DataTypes {
			if item.DataType == oldItem.DataType {
				// Found it, compare & stop
				if item.Count != oldItem.Count {
					if item.DataType == protos.ScanDataType_SD_IMAGE {
						return "image", nil
					} else if item.DataType == protos.ScanDataType_SD_XRF {
						return "spectra", nil
					} else if item.DataType == protos.ScanDataType_SD_RGBU {
						return "rgbu", nil
					}
				}
				break
			}
		}
	}

	compareCounts := []string{"MaxSpectra", "BulkSpectra", "DwellSpectra", "NormalSpectra"}
	for _, name := range compareCounts {
		if newSummary.ContentCounts[name] != oldSummary.ContentCounts[name] {
			return "spectra", nil
		}
	}

	compareCounts = []string{"DriveID", "Site", "Target"}
	for _, name := range compareCounts {
		if newSummary.Meta[name] != oldSummary.Meta[name] {
			return "housekeeping", nil
		}
	}

	if newSummary.Title != oldSummary.Title {
		return "housekeeping", nil
	}

	return "unknown", nil
}