/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF 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
 * <p>
 * http://www.apache.org/licenses/LICENSE-2.0
 * <p>
 * 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.
 */

package org.apache.atlas.repository.migration;

import org.apache.atlas.ApplicationProperties;
import org.apache.atlas.AtlasException;
import org.apache.atlas.RequestContext;
import org.apache.atlas.exception.AtlasBaseException;
import org.apache.atlas.model.impexp.AtlasImportRequest;
import org.apache.atlas.model.migration.MigrationImportStatus;
import org.apache.atlas.repository.graph.AtlasGraphProvider;
import org.apache.atlas.repository.impexp.ImportService;
import org.apache.atlas.type.AtlasType;
import org.apache.commons.codec.digest.DigestUtils;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.io.filefilter.WildcardFileFilter;
import org.apache.commons.lang3.ArrayUtils;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.FileFilter;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.net.InetAddress;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.zip.ZipFile;

public class ZipFileMigrationImporter implements Runnable {
    private static final Logger LOG = LoggerFactory.getLogger(ZipFileMigrationImporter.class);

    private static final String APPLICATION_PROPERTY_MIGRATION_NUMER_OF_WORKERS = "atlas.migration.mode.workers";
    private static final String APPLICATION_PROPERTY_MIGRATION_BATCH_SIZE       = "atlas.migration.mode.batch.size";
    private static final String DEFAULT_NUMBER_OF_WORKERS                       = "4";
    private static final String DEFAULT_BATCH_SIZE                              = "100";
    private static final String ZIP_FILE_COMMENT_ENTITIES_COUNT                 = "entitiesCount";
    private static final String ZIP_FILE_COMMENT_TOTAL_COUNT                    = "total";
    private static final String FILE_EXTENSION_ZIP                              = ".zip";
    private static final String ENV_USER_NAME                                   = "user.name";
    private static final String ARCHIVE_DIR                                     = "archive";

    private final ImportService              importService;
    private       List<String>               filesToImport;
    private final DataMigrationStatusService dataMigrationStatusService;
    private       File                       archiveDir;

    /**
     * Input:
     * fileName : can support wildcards. If it contains wildcards then all matching files will be imported
     */
    public ZipFileMigrationImporter(ImportService importService, String fileName) {
        this.importService              = importService;
        this.dataMigrationStatusService = new DataMigrationStatusService(AtlasGraphProvider.getGraphInstance());

        initialize(fileName);
    }

    @Override
    public void run() {
        for (String fileToImport : filesToImport) {
            try {
                detectFileToImport(fileToImport);

                int streamSize = getStreamSizeFromComment(fileToImport);

                MigrationImportStatus migrationImportStatus = getCreateMigrationStatus(fileToImport, streamSize);

                performImport(fileToImport, streamSize, Long.toString(migrationImportStatus.getCurrentIndex()));

                dataMigrationStatusService.setStatus("DONE");

                moveZipFileToArchiveDir(fileToImport);
            } catch (IOException e) {
                LOG.error("Migration Import: IO Error!", e);

                dataMigrationStatusService.setStatus("FAIL");
            } catch (AtlasBaseException e) {
                LOG.error("Migration Import: Error!", e);

                dataMigrationStatusService.setStatus("FAIL");
            }
        }
    }

    private void initialize(String fileName) {
        this.filesToImport = getAllFilesToImport(fileName);

        if (CollectionUtils.isNotEmpty(this.filesToImport)) {
            createArchiveDirectory(fileName);
        }
    }

    /**
     * Input:
     * fileName : If it contains wildcards then all matching files will be discovered
     */
    private List<String> getAllFilesToImport(String fileName) {
        List<String> ret          = new ArrayList<>();
        File         fileToImport = new File(fileName);

        if (fileToImport.exists() && fileToImport.isFile()) {
            //Input file present so no need to expand
            LOG.info("Migration Import: zip file for import: {}", fileToImport);

            ret.add(fileToImport.getAbsolutePath());
        } else {
            //The fileName might have wildcard
            String dirPath       = new File(fileToImport.getParent()).getAbsolutePath();
            File   importDataDir = new File(dirPath);

            if (importDataDir.exists() && importDataDir.isDirectory()) {
                String     fileNameWithWildcard = fileToImport.getName();
                FileFilter fileFilter           = new WildcardFileFilter(fileNameWithWildcard);
                File[]     importFiles          = importDataDir.listFiles(fileFilter);

                if (ArrayUtils.isNotEmpty(importFiles)) {
                    Arrays.sort(importFiles);

                    LOG.info("Migration Import: zip files for import: ");

                    for (File importFile : importFiles) {
                        if (isValidImportFile(importFile)) {
                            LOG.info("{} with absolute path - {}", importFile.getName(), importFile.getAbsolutePath());

                            ret.add(importFile.getAbsolutePath());
                        } else {
                            LOG.warn("Ignoring {} as it is not a file or does not end with extension {}", FILE_EXTENSION_ZIP, importFile.getAbsolutePath());
                        }
                    }
                } else {
                    LOG.warn("Migration Import: No files to import");
                }
            }
        }

        return ret;
    }

    private boolean isValidImportFile(File importFile) {
        return importFile.isFile() && StringUtils.endsWithIgnoreCase(importFile.getName(), FILE_EXTENSION_ZIP);
    }

    private void createArchiveDirectory(String fileName) {
        File   fileToImport = new File(fileName);
        String parentPath   = new File(fileToImport.getParent()).getAbsolutePath();

        this.archiveDir = new File(parentPath + File.separator + ARCHIVE_DIR);

        if (this.archiveDir.exists() && !this.archiveDir.canWrite()) {
            LOG.warn("Migration Import: No write permission to archive directory {}", this.archiveDir.getAbsolutePath());

            this.archiveDir = null;
        } else if (!this.archiveDir.exists() && !this.archiveDir.getParentFile().canWrite()) {
            LOG.warn("Migration Import: No permission to create archive directory {}", this.archiveDir.getAbsolutePath());

            this.archiveDir = null;
        } else {
            this.archiveDir.mkdirs();

            LOG.info("Migration Import: archive directory for zip files: {}", this.archiveDir.getAbsolutePath());
        }
    }

    private void moveZipFileToArchiveDir(String srcFilePath) {
        if (this.archiveDir == null) {
            return;
        }

        File   sourceFile = new File(srcFilePath);
        String newFile    = archiveDir.getAbsolutePath() + File.separator + sourceFile.getName();

        if (!sourceFile.canWrite()) {
            LOG.warn("Migration Import: No permission to archive the zip file {}", sourceFile.getAbsolutePath());

            this.archiveDir = null;
        } else {
            if (sourceFile.renameTo(new File(newFile))) {
                sourceFile.delete();

                LOG.info("Migration Import: Successfully archived the zip file: {} to {}", srcFilePath, this.archiveDir.getAbsolutePath());
            } else {
                LOG.warn("Migration Import: Failed to archive the zip file: {}", srcFilePath);
            }
        }
    }

    private MigrationImportStatus getCreateMigrationStatus(String fileName, int streamSize) {
        MigrationImportStatus status = null;

        try {
            status = new MigrationImportStatus(fileName, DigestUtils.md5Hex(new FileInputStream(fileName)));

            status.setTotalCount(streamSize);
        } catch (IOException e) {
            LOG.error("Exception occurred while creating migration import", e);
        }

        MigrationImportStatus statusRetrieved = dataMigrationStatusService.getCreate(status);

        LOG.info("DataMigrationStatusService: Position: {}", statusRetrieved.getCurrentIndex());

        dataMigrationStatusService.setStatus("STARTED");

        return statusRetrieved;
    }

    private void detectFileToImport(String fileToImport) throws IOException {
        FileWatcher fileWatcher = new FileWatcher(fileToImport);

        fileWatcher.start();
    }

    private int getStreamSizeFromComment(String fileToImport) {
        int ret = 1;

        try {
            ZipFile zipFile = new ZipFile(fileToImport);
            String  comment = zipFile.getComment();

            ret = processZipFileStreamSizeComment(comment);

            zipFile.close();
        } catch (IOException e) {
            LOG.error("Error opening ZIP file: {}", fileToImport, e);
        }

        return ret;
    }

    private int processZipFileStreamSizeComment(String comment) {
        if (StringUtils.isEmpty(comment)) {
            return 1;
        }

        Map map           = AtlasType.fromJson(comment, Map.class);
        int entitiesCount = (int) map.get(ZIP_FILE_COMMENT_ENTITIES_COUNT);
        int totalCount    = (int) map.get(ZIP_FILE_COMMENT_TOTAL_COUNT);

        LOG.info("ZipFileMigrationImporter: Zip file: Comment: streamSize: {}: total: {}", entitiesCount, totalCount);

        return entitiesCount;
    }

    private void performImport(String fileToImport, int streamSize, String startPosition) throws AtlasBaseException {
        try {
            LOG.info("Migration Import: {}: Starting at: {}...", fileToImport, startPosition);

            InputStream fs = new FileInputStream(fileToImport);

            RequestContext.get().setUser(getUserNameFromEnvironment(), null);

            importService.run(fs, getImportRequest(fileToImport, streamSize, startPosition),
                    getUserNameFromEnvironment(), InetAddress.getLocalHost().getHostName(), InetAddress.getLocalHost().getHostAddress());
        } catch (Exception ex) {
            LOG.error("Migration Import: Error loading zip for migration!", ex);

            throw new AtlasBaseException(ex);
        } finally {
            LOG.info("Migration Import: {}: Done!", fileToImport);
        }
    }

    private String getUserNameFromEnvironment() {
        return System.getProperty(ENV_USER_NAME);
    }

    private AtlasImportRequest getImportRequest(String fileToImport, int streamSize, String position) throws AtlasException {
        AtlasImportRequest request = new AtlasImportRequest();

        request.setOption(AtlasImportRequest.OPTION_KEY_MIGRATION_FILE_NAME, fileToImport);
        request.setSizeOption(streamSize);
        request.setOption(AtlasImportRequest.OPTION_KEY_MIGRATION, "true");
        request.setOption(AtlasImportRequest.OPTION_KEY_NUM_WORKERS, getPropertyValue(APPLICATION_PROPERTY_MIGRATION_NUMER_OF_WORKERS, DEFAULT_NUMBER_OF_WORKERS));
        request.setOption(AtlasImportRequest.OPTION_KEY_BATCH_SIZE, getPropertyValue(APPLICATION_PROPERTY_MIGRATION_BATCH_SIZE, DEFAULT_BATCH_SIZE));
        request.setOption(AtlasImportRequest.START_POSITION_KEY, (StringUtils.isEmpty(position) ? "0" : position));

        return request;
    }

    private String getPropertyValue(String property, String defaultValue) throws AtlasException {
        return ApplicationProperties.get().getString(property, defaultValue);
    }
}
