/**
 * 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.store.graph.v2.tasks.searchdownload;

import com.opencsv.CSVWriter;
import org.apache.atlas.ApplicationProperties;
import org.apache.atlas.AtlasConfiguration;
import org.apache.atlas.AtlasException;
import org.apache.atlas.RequestContext;
import org.apache.atlas.discovery.AtlasDiscoveryService;
import org.apache.atlas.exception.AtlasBaseException;
import org.apache.atlas.model.discovery.AtlasSearchResult;
import org.apache.atlas.model.discovery.SearchParameters;
import org.apache.atlas.model.instance.AtlasEntityHeader;
import org.apache.atlas.model.instance.AtlasObjectId;
import org.apache.atlas.model.tasks.AtlasTask;
import org.apache.atlas.tasks.AbstractTask;
import org.apache.atlas.type.AtlasTypeRegistry;
import org.apache.atlas.utils.AtlasJson;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.collections4.MapUtils;
import org.apache.commons.configuration.Configuration;
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.FileWriter;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import static org.apache.atlas.model.discovery.AtlasSearchResult.AtlasQueryType.BASIC;
import static org.apache.atlas.model.discovery.AtlasSearchResult.AtlasQueryType.DSL;
import static org.apache.atlas.model.tasks.AtlasTask.Status.COMPLETE;
import static org.apache.atlas.model.tasks.AtlasTask.Status.FAILED;

public class SearchResultDownloadTask extends AbstractTask {
    private static final Logger LOG = LoggerFactory.getLogger(SearchResultDownloadTask.class);

    public static final String SEARCH_PARAMETERS_JSON_KEY = "search_parameters_json";
    public static final String CSV_FILE_NAME_KEY          = "csv_file_Name";
    public static final String SEARCH_TYPE_KEY            = "search_type";
    public static final String ATTRIBUTE_LABEL_MAP_KEY    = "attribute_label_map";
    public static final String QUERY_KEY                  = "query";
    public static final String TYPE_NAME_KEY              = "type_name";
    public static final String CLASSIFICATION_KEY         = "classification";
    public static final String LIMIT_KEY                  = "limit";
    public static final String OFFSET_KEY                 = "offset";
    public static final String CSV_FILE_EXTENSION         = ".csv";
    public static final String DOWNLOAD_DIR_PATH;

    private static final String EMPTY_STRING               = "";
    private static final String DOWNLOAD_DIR_PATH_KEY      = "atlas.download.search.dir.path";
    private static final String DOWNLOAD_DIR_PATH_DEFAULT  = System.getProperty("user.dir");
    private static final String CSV_DOWNLOAD_DIR           = "search_result_downloads";

    private final AtlasDiscoveryService discoveryService;
    private final AtlasTypeRegistry     typeRegistry;

    public SearchResultDownloadTask(AtlasTask task, AtlasDiscoveryService discoveryService, AtlasTypeRegistry typeRegistry) {
        super(task);

        this.discoveryService = discoveryService;
        this.typeRegistry     = typeRegistry;
    }

    @Override
    public AtlasTask.Status perform() throws Exception {
        RequestContext.clear();

        Map<String, Object> params = getTaskDef().getParameters();

        if (MapUtils.isEmpty(params)) {
            LOG.warn("Task: {}: Unable to process task: Parameters is not readable!", getTaskGuid());

            return FAILED;
        }

        String userName = getTaskDef().getCreatedBy();

        if (StringUtils.isEmpty(userName)) {
            LOG.warn("Task: {}: Unable to process task as user name is empty!", getTaskGuid());

            return FAILED;
        }

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

        try {
            run(params);

            setStatus(COMPLETE);
        } catch (Exception e) {
            LOG.error("Task: {}: Error performing task!", getTaskGuid(), e);

            setStatus(FAILED);

            throw e;
        } finally {
            RequestContext.clear();
        }

        return getStatus();
    }

    protected void run(Map<String, Object> parameters) throws AtlasBaseException, IOException {
        Map<String, String>              attributeLabelMap;
        AtlasSearchResult                searchResult = null;
        AtlasSearchResult.AtlasQueryType queryType    = null;

        if (parameters.get(SEARCH_TYPE_KEY) == BASIC) {
            String           searchParametersJson = (String) parameters.get(SEARCH_PARAMETERS_JSON_KEY);
            SearchParameters searchParameters     = AtlasJson.fromJson(searchParametersJson, SearchParameters.class);

            searchParameters.setLimit(AtlasConfiguration.SEARCH_MAX_LIMIT.getInt());

            searchResult = discoveryService.searchWithParameters(searchParameters);
            queryType    = BASIC;
        } else if (parameters.get(SEARCH_TYPE_KEY) == DSL) {
            String query          = (String) parameters.get(QUERY_KEY);
            String typeName       = (String) parameters.get(TYPE_NAME_KEY);
            String classification = (String) parameters.get(CLASSIFICATION_KEY);
            int    offset         = (int) parameters.get(OFFSET_KEY);
            String queryStr       = discoveryService.getDslQueryUsingTypeNameClassification(query, typeName, classification);

            searchResult = discoveryService.searchUsingDslQuery(queryStr, AtlasConfiguration.SEARCH_MAX_LIMIT.getInt(), offset);
            queryType    = DSL;
        }

        String attributeLabelMapJson = (String) parameters.get(ATTRIBUTE_LABEL_MAP_KEY);

        attributeLabelMap = AtlasJson.fromJson(attributeLabelMapJson, Map.class);

        if (searchResult != null) {
            generateCSVFileFromSearchResult(searchResult, attributeLabelMap, queryType);
        }
    }

    private void generateCSVFileFromSearchResult(AtlasSearchResult searchResult, Map<String, String> attributeLabelMap, AtlasSearchResult.AtlasQueryType queryType) throws IOException {
        List<AtlasEntityHeader>                 allEntityHeaders      = searchResult.getEntities();
        AtlasSearchResult.AttributeSearchResult attributeSearchResult = searchResult.getAttributes();
        String                                  fileName              = (String) getTaskDef().getParameters().get(CSV_FILE_NAME_KEY);

        if ((queryType == BASIC && CollectionUtils.isEmpty(allEntityHeaders)) || (queryType == DSL && (CollectionUtils.isEmpty(allEntityHeaders) && attributeSearchResult == null))) {
            LOG.info("No result found. Not generating csv file: {}", fileName);

            return;
        }

        File dir = new File(DOWNLOAD_DIR_PATH, RequestContext.getCurrentUser());

        if (!dir.exists()) {
            dir.mkdirs();
        }

        File csvFile = new File(dir, fileName);

        try (FileWriter fileWriter = new FileWriter(csvFile); CSVWriter csvWriter = new CSVWriter(fileWriter)) {
            String[] defaultHeaders = new String[] {"Type name", "Name", "Classifications", "Terms"};
            String[] attributeHeaders;
            int      attrSize;

            if (attributeLabelMap == null) {
                attributeLabelMap = new HashMap<>();
            }

            attributeLabelMap.put("Owner", "owner");
            attributeLabelMap.put("Description", "description");

            Collection<String> attributeHeaderLabels = attributeLabelMap.keySet();

            if (queryType == DSL && (CollectionUtils.isEmpty(allEntityHeaders) && attributeSearchResult != null)) {
                attributeHeaderLabels = attributeSearchResult.getName();
                defaultHeaders        = new String[0];
            }

            attrSize         = (attributeHeaderLabels == null) ? 0 : attributeHeaderLabels.size();
            attributeHeaders = new String[attrSize];

            if (attributeHeaderLabels != null) {
                attributeHeaders = attributeHeaderLabels.toArray(attributeHeaders);
            }

            int      headerSize = attrSize + defaultHeaders.length;
            String[] headers    = new String[headerSize];

            System.arraycopy(defaultHeaders, 0, headers, 0, defaultHeaders.length);

            if (ArrayUtils.isNotEmpty(attributeHeaders)) {
                System.arraycopy(attributeHeaders, 0, headers, defaultHeaders.length, attrSize);
            }

            csvWriter.writeNext(headers);

            String[] entityRecords = new String[headerSize];

            if (CollectionUtils.isNotEmpty(allEntityHeaders)) {
                for (AtlasEntityHeader entityHeader : allEntityHeaders) {
                    entityRecords[0] = entityHeader.getTypeName();
                    entityRecords[1] = entityHeader.getDisplayText() != null ? entityHeader.getDisplayText() : entityHeader.getGuid();
                    entityRecords[2] = String.join(",", entityHeader.getClassificationNames());
                    entityRecords[3] = String.join(",", entityHeader.getMeaningNames());

                    if (MapUtils.isNotEmpty(entityHeader.getAttributes())) {
                        for (int i = defaultHeaders.length; i < headerSize; i++) {
                            Object attrValue = entityHeader.getAttribute(attributeLabelMap.get(headers[i]));

                            if (attrValue instanceof AtlasObjectId) {
                                entityRecords[i] = String.valueOf(((AtlasObjectId) attrValue).getUniqueAttributes().get("qualifiedName"));
                            } else if (attrValue instanceof List) {
                                if (CollectionUtils.isNotEmpty((List<?>) attrValue)) {
                                    List<String> valueList = new ArrayList<>();

                                    for (Object attrVal : (List<?>) attrValue) {
                                        if (attrVal instanceof AtlasObjectId) {
                                            String value = String.valueOf(((AtlasObjectId) attrVal).getUniqueAttributes().get("qualifiedName"));

                                            valueList.add(value);
                                        } else {
                                            valueList.add(String.valueOf(attrVal));
                                        }
                                    }

                                    entityRecords[i] = String.join(",", valueList);
                                }
                            } else {
                                entityRecords[i] = attrValue == null ? EMPTY_STRING : String.valueOf(attrValue);
                            }
                        }
                    }

                    csvWriter.writeNext(entityRecords);
                }
            }

            if (queryType == DSL && attributeSearchResult != null) {
                for (List<Object> resultSet : attributeSearchResult.getValues()) {
                    for (int i = defaultHeaders.length; i < headerSize; i++) {
                        entityRecords[i] = resultSet.get(i) == null ? EMPTY_STRING : String.valueOf(resultSet.get(i));
                    }

                    csvWriter.writeNext(entityRecords);
                }
            }
        }
    }

    static {
        Configuration configuration = null;

        try {
            configuration = ApplicationProperties.get();
        } catch (AtlasException e) {
            LOG.error("Failed to load application properties", e);
        }

        if (configuration != null) {
            DOWNLOAD_DIR_PATH = configuration.getString(DOWNLOAD_DIR_PATH_KEY, DOWNLOAD_DIR_PATH_DEFAULT) + File.separator + CSV_DOWNLOAD_DIR;
        } else {
            DOWNLOAD_DIR_PATH = DOWNLOAD_DIR_PATH_DEFAULT + File.separator + CSV_DOWNLOAD_DIR;
        }
    }
}
