/**
 * 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.tools;

import com.sun.jersey.core.util.MultivaluedMapImpl;
import org.apache.atlas.ApplicationProperties;
import org.apache.atlas.AtlasClientV2;
import org.apache.atlas.AtlasException;
import org.apache.atlas.AtlasServiceException;
import org.apache.atlas.model.SearchFilter;
import org.apache.atlas.model.instance.AtlasClassification;
import org.apache.atlas.model.instance.AtlasEntityHeader;
import org.apache.atlas.model.instance.AtlasEntityHeaders;
import org.apache.atlas.model.typedef.AtlasClassificationDef;
import org.apache.atlas.model.typedef.AtlasEntityDef;
import org.apache.atlas.model.typedef.AtlasStructDef;
import org.apache.atlas.model.typedef.AtlasTypesDef;
import org.apache.atlas.type.AtlasType;
import org.apache.atlas.utils.AtlasJson;
import org.apache.atlas.utils.AuthenticationUtil;
import org.apache.commons.cli.CommandLine;
import org.apache.commons.cli.DefaultParser;
import org.apache.commons.cli.Options;
import org.apache.commons.cli.ParseException;
import org.apache.commons.configuration.Configuration;
import org.apache.commons.lang3.StringUtils;
import org.apache.hadoop.security.UserGroupInformation;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.ws.rs.core.MultivaluedMap;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileReader;
import java.io.FileWriter;
import java.io.IOException;
import java.text.SimpleDateFormat;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.TimeZone;
import java.util.function.Consumer;

import static org.apache.atlas.model.instance.AtlasEntity.Status.DELETED;

public class BulkFetchAndUpdate {
    private static final Logger LOG = LoggerFactory.getLogger(BulkFetchAndUpdate.class);

    private static final String DATE_FORMAT_SUPPORTED               = "yyyy-MM-dd'T'HH:mm:ss";
    private static final String OPTION_FROM                         = "f";
    private static final String APPLICATION_PROPERTY_ATLAS_ENDPOINT = "atlas.rest.address";
    private static final String SYSTEM_PROPERTY_USER_DIR            = "user.dir";
    private static final String STEP_PREPARE                        = "prepare";
    private static final String STEP_UPDATE                         = "update";
    private static final String DEFAULT_ATLAS_URL                   = "http://localhost:21000/";
    private static final String FILE_CLASSIFICATION_DEFS            = "classification-definitions.json";
    private static final String FILE_ENTITY_HEADERS                 = "entity-headers.json";
    private static final int    EXIT_CODE_SUCCESS                   = 0;
    private static final int    EXIT_CODE_FAILED                    = 1;

    private BulkFetchAndUpdate() {
        // to block instantiation
    }

    private static final String[] filesToUse = new String[] {
            FILE_CLASSIFICATION_DEFS,
            FILE_ENTITY_HEADERS
    };

    public static void main(String[] args) {
        int exitCode = EXIT_CODE_FAILED;

        try {
            long        fromTimestamp = 0L;
            CommandLine cmd           = getCommandLine(args);

            String stepToExecute = cmd.getOptionValue("s").trim();
            String uid           = cmd.getOptionValue("u");
            String pwd           = cmd.getOptionValue("p");
            String directory     = cmd.getOptionValue("d");
            String fromTime      = cmd.getOptionValue(OPTION_FROM);
            String basePath      = getDirectory(directory);

            displayCrLf(basePath);

            String[] atlasEndpoint = getAtlasRESTUrl();
            if (atlasEndpoint == null || atlasEndpoint.length == 0) {
                atlasEndpoint = new String[] {DEFAULT_ATLAS_URL};
            }

            if (StringUtils.equals(stepToExecute, STEP_PREPARE)) {
                if (StringUtils.isEmpty(fromTime)) {
                    displayCrLf("'fromTime' is empty" + fromTime);
                    printUsage();
                    return;
                }

                fromTimestamp = getTimestamp(fromTime);
                displayCrLf("fromTimestamp: " + fromTimestamp);
                if (fromTimestamp == 0L) {
                    printUsage();
                    return;
                }
            }

            process(stepToExecute, basePath, atlasEndpoint, uid, pwd, fromTimestamp);

            exitCode = EXIT_CODE_SUCCESS;
        } catch (ParseException e) {
            LOG.error("Failed to parse arguments. Error: ", e.getMessage());
            printUsage();
        } catch (Exception e) {
            LOG.error("Failed!", e);
            displayCrLf("Failed: " + e.getMessage());
        }

        System.exit(exitCode);
    }

    private static long getTimestamp(String str) {
        try {
            if (StringUtils.isEmpty(str)) {
                return 0;
            }

            TimeZone         utc              = TimeZone.getDefault();
            SimpleDateFormat simpleDateFormat = new SimpleDateFormat(DATE_FORMAT_SUPPORTED);
            simpleDateFormat.setTimeZone(utc);

            return simpleDateFormat.parse(str).getTime();
        } catch (java.text.ParseException e) {
            displayCrLf("Unsupported date format: " + str);
            return 0;
        }
    }

    private static void process(String stepToExecute, String basePath, String[] atlasEndpoint, String uid, String pwd, long fromTimestamp) throws Exception {
        AtlasClientV2 atlasClientV2 = getAtlasClientV2(atlasEndpoint, new String[] {uid, pwd});

        switch (stepToExecute) {
            case STEP_PREPARE: {
                Preparer p = new Preparer(atlasClientV2);
                p.run(basePath, fromTimestamp);
            }
            break;

            case STEP_UPDATE: {
                Updater u = new Updater(atlasClientV2);
                u.run(basePath);
            }
            break;

            default:
                printUsage();
                break;
        }
    }

    private static String getDirectory(String directory) {
        String basePath = System.getProperty(SYSTEM_PROPERTY_USER_DIR) + File.separatorChar;
        if (StringUtils.isNotEmpty(directory) && checkDirectoryExists(directory)) {
            basePath = directory + File.separatorChar;
        } else {
            display("Using directory: ");
        }

        return basePath;
    }

    private static CommandLine getCommandLine(String[] args) throws ParseException {
        Options options = new Options();
        options.addRequiredOption("s", "step", true, "Step to run.");
        options.addOption("u", "user", true, "User name.");
        options.addOption("p", "password", true, "Password name.");
        options.addOption("d", "dir", true, "Directory for reading/writing data.");
        options.addOption(OPTION_FROM, "fromDate", true, "Date, in YYYY-MM-DD format, from where to start reading.");

        return new DefaultParser().parse(options, args);
    }

    private static void printUsage() {
        System.out.println();
        displayCrLf("Usage: classification-updater.sh [-s <step>] [-f <from time>] [-t <optional: to time>] [-d <dir>]");
        displayCrLf("    step: Specify which step to execute:");
        displayCrLf("           prepare: prepare classifications and associated entities.");
        displayCrLf("           update: update classifications and entities.");
        displayCrLf("    dir: [optional] Directory where read/write will happen.");
        displayCrLf("           If not specified, current directory will be used.");
        displayCrLf("    from: [mandatory for 'prepare' step, optional for 'update' step] Date, in YYYY-MM-DD format, from where audits need to be read.");
        displayCrLf("           If not specified, current directory will be used.");
        System.out.println();
    }

    private static String[] getAtlasRESTUrl() {
        Configuration atlasConf;
        try {
            atlasConf = ApplicationProperties.get();
            return atlasConf.getStringArray(APPLICATION_PROPERTY_ATLAS_ENDPOINT);
        } catch (AtlasException e) {
            return new String[] {DEFAULT_ATLAS_URL};
        }
    }

    private static AtlasClientV2 getAtlasClientV2(String[] atlasEndpoint, String[] uidPwdFromCommandLine) throws IOException {
        AtlasClientV2 atlasClientV2;
        if (!AuthenticationUtil.isKerberosAuthenticationEnabled()) {
            String[] uidPwd = (uidPwdFromCommandLine[0] == null || uidPwdFromCommandLine[1] == null) ? AuthenticationUtil.getBasicAuthenticationInput() : uidPwdFromCommandLine;

            atlasClientV2 = new AtlasClientV2(atlasEndpoint, uidPwd);
        } else {
            UserGroupInformation ugi = UserGroupInformation.getCurrentUser();
            atlasClientV2 = new AtlasClientV2(ugi, ugi.getShortUserName(), atlasEndpoint);
        }
        return atlasClientV2;
    }

    private static void displayCrLf(String... formatMessage) {
        displayFn(System.out::println, formatMessage);
    }

    private static void display(String... formatMessage) {
        displayFn(System.out::print, formatMessage);
    }

    private static void displayFn(Consumer<String> fn, String... formatMessage) {
        if (formatMessage.length == 1) {
            fn.accept(formatMessage[0]);
        } else {
            fn.accept(String.format(formatMessage[0], formatMessage[1]));
        }
    }

    private static void closeReader(BufferedReader bufferedReader) {
        try {
            if (bufferedReader == null) {
                return;
            }

            bufferedReader.close();
        } catch (IOException ex) {
            LOG.error("closeReader", ex);
        }
    }

    private static BufferedReader getBufferedReader(String basePath, String fileName) throws FileNotFoundException {
        return new BufferedReader(new FileReader(basePath + fileName));
    }

    private static boolean fileCheck(String basePath, String[] files, boolean existCheck) {
        boolean ret = true;
        for (String f : files) {
            ret = ret && fileCheck(basePath, f, existCheck);
        }

        return ret;
    }

    private static boolean fileCheck(String basePath, String file, boolean existCheck) {
        String errorMessage = existCheck ? "does not exist" : "exists";
        if (checkFileExists(basePath + file) != existCheck) {
            displayCrLf(String.format("File '%s' %s!", basePath + file, errorMessage));
            return false;
        }

        return true;
    }

    private static boolean checkFileExists(String fileName) {
        File f = new File(fileName);
        return f.exists() && !f.isDirectory();
    }

    private static boolean checkDirectoryExists(String fileName) {
        File f = new File(fileName);
        return f.exists() && f.isDirectory();
    }

    private static FileWriter getFileWriter(String basePath, String fileName) throws IOException {
        String filePath = basePath + fileName;
        displayCrLf("Creating %s", filePath);
        return new FileWriter(filePath, true);
    }

    private static class Preparer {
        private static final String ATTR_NAME_QUALIFIED_NAME = "qualifiedName";

        private final AtlasClientV2 atlasClientV2;
        private final Map<String, String> typeNameUniqueAttributeNameMap = new HashMap<>();

        public Preparer(AtlasClientV2 atlasClientV2) {
            this.atlasClientV2 = atlasClientV2;
        }

        public void run(String basePath, long fromTimestamp) throws Exception {
            if (!fileCheck(basePath, filesToUse, false)) {
                return;
            }

            displayCrLf("Starting: from: " + fromTimestamp + " to: " + "current time (" + System.currentTimeMillis() + ")...");
            writeClassificationDefs(basePath, FILE_CLASSIFICATION_DEFS, getAllClassificationsDefs());
            writeEntityHeaders(basePath, FILE_ENTITY_HEADERS, fromTimestamp);
            displayCrLf("Done!");
        }

        private void writeClassificationDefs(String basePath, String fileName, List<AtlasClassificationDef> classificationDefs) throws IOException {
            try (FileWriter fileWriter = getFileWriter(basePath, fileName)) {
                for (AtlasClassificationDef classificationDef : classificationDefs) {
                    try {
                        classificationDef.setGuid(null);
                        String json = AtlasType.toJson(classificationDef);
                        fileWriter.write(json + "\n");
                    } catch (Exception e) {
                        LOG.error("Error writing classifications: {}", e);
                        displayCrLf("Error writing classifications.");
                    }
                }
            }
        }

        private void writeEntityHeaders(String basePath, String fileName, long fromTimestamp) throws AtlasServiceException, IOException {
            FileWriter fileWriter;

            try {
                fileWriter = getFileWriter(basePath, fileName);
            } catch (IOException e) {
                LOG.error("Error opening {}/{}", basePath, fileName, e);
                displayCrLf("Error opening: %", basePath + File.separatorChar + fileName);
                return;
            }

            try {
                AtlasEntityHeaders entityHeaders     = atlasClientV2.getEntityHeaders(fromTimestamp);
                int                guidHeaderMapSize = entityHeaders.getGuidHeaderMap().size();
                try {
                    displayCrLf("Read entities: " + guidHeaderMapSize);
                    AtlasEntityHeaders updatedHeaders = removeEntityGuids(entityHeaders);
                    fileWriter.write(AtlasType.toJson(updatedHeaders));

                    displayCrLf("Writing entities: " + updatedHeaders.getGuidHeaderMap().size());
                } catch (Exception e) {
                    LOG.error("Error writing: {}", guidHeaderMapSize, e);
                    displayCrLf("Error writing: " + e);
                }
            } finally {
                if (fileWriter != null) {
                    fileWriter.close();
                }
            }
        }

        private AtlasEntityHeaders removeEntityGuids(AtlasEntityHeaders headers) {
            Map<String, AtlasEntityHeader> uniqueNameEntityHeaderMap = new HashMap<>();

            for (AtlasEntityHeader header : headers.getGuidHeaderMap().values()) {
                String uniqueName = getUniqueName(header);
                if (StringUtils.isEmpty(uniqueName)) {
                    displayCrLf("UniqueName is empty.  Ignoring: " + header.getGuid());
                    LOG.warn("UniqueName is empty. Ignoring: {}", AtlasJson.toJson(header));
                    continue;
                }

                displayCrLf("Processing: " + uniqueName);
                if (header.getStatus() == DELETED) {
                    continue;
                }

                String key = String.format("%s:%s", header.getTypeName(), uniqueName);
                header.setGuid(null);
                boolean keyFound = uniqueNameEntityHeaderMap.containsKey(key);
                if (!keyFound) {
                    uniqueNameEntityHeaderMap.put(key, header);
                }

                updateClassificationsForHeader(header, uniqueNameEntityHeaderMap.get(key), keyFound);
                displayCrLf("Processing: " + uniqueName);
            }

            displayCrLf("Processed: " + uniqueNameEntityHeaderMap.size());
            headers.setGuidHeaderMap(uniqueNameEntityHeaderMap);
            return headers;
        }

        private void updateClassificationsForHeader(AtlasEntityHeader header, AtlasEntityHeader currentHeader, boolean keyFound) {
            for (AtlasClassification c : header.getClassifications()) {
                c.setEntityGuid(null);

                if (keyFound) {
                    boolean found =
                            currentHeader.getClassifications().stream().anyMatch(ox -> ox.getTypeName().equals(c.getTypeName()));
                    if (!found) {
                        currentHeader.getClassifications().add(c);
                    } else {
                        displayCrLf("Ignoring: " + c);
                        LOG.warn("Ignoring: {}", AtlasJson.toJson(c));
                    }
                }
            }
        }

        private String getUniqueName(AtlasEntityHeader header) {
            String uniqueAttributeName = ATTR_NAME_QUALIFIED_NAME;
            if (!header.getAttributes().containsKey(ATTR_NAME_QUALIFIED_NAME)) {
                uniqueAttributeName = getUniqueAttribute(header.getTypeName());
            }

            Object attrValue = header.getAttribute(uniqueAttributeName);
            if (attrValue == null) {
                LOG.warn("Unique Attribute Value: empty: {}", AtlasJson.toJson(header));
                return StringUtils.EMPTY;
            }

            return attrValue.toString();
        }

        private String getUniqueAttribute(String typeName) {
            try {
                if (typeNameUniqueAttributeNameMap.containsKey(typeName)) {
                    return typeNameUniqueAttributeNameMap.get(typeName);
                }

                AtlasEntityDef entityDef = atlasClientV2.getEntityDefByName(typeName);
                for (AtlasStructDef.AtlasAttributeDef ad : entityDef.getAttributeDefs()) {
                    if (ad.getIsUnique()) {
                        typeNameUniqueAttributeNameMap.put(typeName, ad.getName());
                        return ad.getName();
                    }
                }
            } catch (AtlasServiceException e) {
                LOG.error("Error fetching type: {}", typeName, e);
                return null;
            }

            return null;
        }

        private List<AtlasClassificationDef> getAllClassificationsDefs() throws Exception {
            MultivaluedMap<String, String> searchParams = new MultivaluedMapImpl();
            searchParams.add(SearchFilter.PARAM_TYPE, "CLASSIFICATION");
            SearchFilter searchFilter = new SearchFilter(searchParams);

            AtlasTypesDef typesDef = atlasClientV2.getAllTypeDefs(searchFilter);
            displayCrLf("Found classifications: " + typesDef.getClassificationDefs().size());
            return typesDef.getClassificationDefs();
        }
    }

    private static class Updater {
        private final AtlasClientV2 atlasClientV2;

        public Updater(AtlasClientV2 atlasClientV2) {
            this.atlasClientV2 = atlasClientV2;
        }

        public void run(String basePath) throws Exception {
            if (!fileCheck(basePath, filesToUse, true)) {
                return;
            }

            displayCrLf("Starting...");
            readAndCreateOrUpdateClassificationDefs(basePath, FILE_CLASSIFICATION_DEFS);
            readEntityUpdates(basePath, FILE_ENTITY_HEADERS);
            displayCrLf("Done!");
        }

        private void readEntityUpdates(String basePath, String fileName) throws IOException {
            BufferedReader bufferedReader = null;
            try {
                bufferedReader = getBufferedReader(basePath, fileName);
                String json = bufferedReader.readLine();
                if (StringUtils.isEmpty(json)) {
                    displayCrLf("Empty file encountered: %s", fileName);
                    return;
                }

                AtlasEntityHeaders response = AtlasType.fromJson(json, AtlasEntityHeaders.class);
                displayCrLf("Found :" + response.getGuidHeaderMap().size());
                String output = atlasClientV2.setClassifications(response);
                displayCrLf(output);
            } catch (AtlasServiceException e) {
                displayCrLf("Error updating. Please see log for details.");
                LOG.error("Error updating. {}", e);
            } finally {
                closeReader(bufferedReader);
            }
        }

        private void readAndCreateOrUpdateClassificationDefs(String basePath, String fileName) throws Exception {
            BufferedReader bufferedReader = null;
            try {
                String cd;
                bufferedReader = getBufferedReader(basePath, fileName);

                while ((cd = bufferedReader.readLine()) != null) {
                    AtlasClassificationDef classificationDef = AtlasType.fromJson(cd, AtlasClassificationDef.class);
                    createOrUpdateClassification(classificationDef);
                }
            } finally {
                closeReader(bufferedReader);
            }
        }

        private void createOrUpdateClassification(AtlasClassificationDef classificationDef) {
            String        name     = classificationDef.getName();
            AtlasTypesDef typesDef = new AtlasTypesDef(null, null, Collections.singletonList(classificationDef), null, null);
            try {
                display("%s -> ", name);
                atlasClientV2.createAtlasTypeDefs(typesDef);
                displayCrLf(" [Done]");
            } catch (AtlasServiceException e) {
                LOG.error("{} skipped!", name, e);
                displayCrLf(" [Skipped]", name);
                updateClassification(classificationDef);
                displayCrLf(" [Done!]");
            } catch (Exception ex) {
                LOG.error("{} skipped!", name, ex);
                displayCrLf(" [Skipped]", name);
            }
        }

        private void updateClassification(AtlasClassificationDef classificationDef) {
            String        name     = classificationDef.getName();
            AtlasTypesDef typesDef = new AtlasTypesDef(null, null, Collections.singletonList(classificationDef), null, null);
            try {
                display("Update: %s -> ", name);
                atlasClientV2.updateAtlasTypeDefs(typesDef);
                displayCrLf(" [Done]");
            } catch (Exception ex) {
                LOG.error("{} skipped!", name, ex);
                displayCrLf(" [Skipped]", name);
            }
        }
    }
}
