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

import org.apache.atlas.AtlasException;
import org.apache.atlas.exception.AtlasBaseException;
import org.apache.commons.cli.BasicParser;
import org.apache.commons.cli.CommandLine;
import org.apache.commons.cli.HelpFormatter;
import org.apache.commons.cli.Options;
import org.apache.commons.cli.ParseException;
import org.apache.commons.lang3.StringUtils;
import org.quartz.CronExpression;
import org.quartz.CronScheduleBuilder;
import org.quartz.DisallowConcurrentExecution;
import org.quartz.Job;
import org.quartz.JobBuilder;
import org.quartz.JobDetail;
import org.quartz.JobExecutionContext;
import org.quartz.JobExecutionException;
import org.quartz.Scheduler;
import org.quartz.SchedulerException;
import org.quartz.Trigger;
import org.quartz.TriggerBuilder;
import org.quartz.impl.StdSchedulerFactory;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.List;

import static org.apache.atlas.trino.cli.ExtractorContext.OPTION_CATALOG_LONG;
import static org.apache.atlas.trino.cli.ExtractorContext.OPTION_CATALOG_SHORT;
import static org.apache.atlas.trino.cli.ExtractorContext.OPTION_CRON_EXPRESSION_LONG;
import static org.apache.atlas.trino.cli.ExtractorContext.OPTION_CRON_EXPRESSION_SHORT;
import static org.apache.atlas.trino.cli.ExtractorContext.OPTION_HELP_LONG;
import static org.apache.atlas.trino.cli.ExtractorContext.OPTION_HELP_SHORT;
import static org.apache.atlas.trino.cli.ExtractorContext.OPTION_SCHEMA_LONG;
import static org.apache.atlas.trino.cli.ExtractorContext.OPTION_SCHEMA_SHORT;
import static org.apache.atlas.trino.cli.ExtractorContext.OPTION_TABLE_LONG;
import static org.apache.atlas.trino.cli.ExtractorContext.OPTION_TABLE_SHORT;

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

    private static final int EXIT_CODE_SUCCESS = 0;
    private static final int EXIT_CODE_FAILED  = 1;
    private static final int EXIT_CODE_HELP    = 2;

    private static TrinoExtractor instance;

    private final ExtractorContext extractorContext;
    private       int              exitCode = EXIT_CODE_FAILED;

    private TrinoExtractor(String[] args) throws Exception {
        extractorContext = createExtractorContext(args);
    }

    public static void main(String[] args) {
        try {
            instance = new TrinoExtractor(args);

            if (instance.extractorContext != null) {
                instance.run();
            } else {
                LOG.error("Extractor context is null. Cannot proceed with extraction.");

                instance.exitCode = EXIT_CODE_FAILED;
            }
        } catch (Exception e) {
            LOG.error("Extraction failed", e);

            instance.exitCode = EXIT_CODE_FAILED;
        }

        System.exit(instance != null ? instance.exitCode : EXIT_CODE_FAILED);
    }

    private void run() {
        try {
            String cronExpression = extractorContext.getCronExpression();

            if (StringUtils.isNotEmpty(cronExpression)) {
                if (!CronExpression.isValidExpression(cronExpression)) {
                    LOG.error("Invalid cron expression: {}", cronExpression);

                    exitCode = EXIT_CODE_FAILED;
                } else {
                    LOG.info("Scheduling extraction for cron expression: {}", cronExpression);

                    JobDetail job = JobBuilder.newJob(MetadataJob.class).withIdentity("metadataJob", "group1").build();
                    Trigger trigger = TriggerBuilder.newTrigger()
                            .withSchedule(CronScheduleBuilder.cronSchedule(cronExpression)).startNow()
                            .build();

                    Scheduler scheduler = null;

                    try {
                        scheduler = new StdSchedulerFactory().getScheduler();

                        scheduler.scheduleJob(job, trigger);
                        scheduler.start();

                        try {
                            Thread.currentThread().join();
                        } catch (InterruptedException ie) {
                            LOG.info("Main thread interrupted, shutting down scheduler.");

                            exitCode = EXIT_CODE_SUCCESS;
                        }
                    } finally {
                        try {
                            if (scheduler != null && !scheduler.isShutdown()) {
                                scheduler.shutdown(true);
                            }
                        } catch (SchedulerException se) {
                            LOG.warn("Error shutting down scheduler: {}", se.getMessage());
                        }
                    }
                }
            } else {
                LOG.info("Running extraction once");

                ExtractorService extractorService = new ExtractorService();

                if (extractorService.execute(extractorContext)) {
                    exitCode = EXIT_CODE_SUCCESS;
                }
            }
        } catch (Exception e) {
            LOG.error("Error encountered. exitCode: {}", exitCode, e);
        } finally {
            if (extractorContext.getAtlasConnector() != null) {
                extractorContext.getAtlasConnector().close();
            }
        }

        System.exit(exitCode);
    }

    private ExtractorContext createExtractorContext(String[] args) throws AtlasBaseException, IOException {
        Options acceptedCliOptions = prepareCommandLineOptions();

        try {
            CommandLine  cmd              = new BasicParser().parse(acceptedCliOptions, args, true);
            List<String> argsNotProcessed = cmd.getArgList();

            if (argsNotProcessed != null && !argsNotProcessed.isEmpty()) {
                throw new AtlasBaseException("Unrecognized arguments.");
            }

            ExtractorContext ret = null;

            if (cmd.hasOption(ExtractorContext.OPTION_HELP_SHORT)) {
                printUsage(acceptedCliOptions);
                exitCode = EXIT_CODE_HELP;
            } else {
                ret = new ExtractorContext(cmd);
                LOG.debug("Successfully initialized the extractor context.");
            }

            return ret;
        } catch (ParseException | AtlasBaseException e) {
            printUsage(acceptedCliOptions);

            throw new AtlasBaseException("Invalid arguments. Reason: " + e.getMessage(), e);
        } catch (AtlasException e) {
            throw new AtlasBaseException("Error in getting Application Properties. Reason: " + e.getMessage(), e);
        } catch (IOException e) {
            throw new IOException(e);
        }
    }

    private static Options prepareCommandLineOptions() {
        Options acceptedCliOptions = new Options();

        return acceptedCliOptions.addOption(OPTION_CATALOG_SHORT, OPTION_CATALOG_LONG, true, "Catalog name")
                .addOption(OPTION_SCHEMA_SHORT, OPTION_SCHEMA_LONG, true, "Schema name")
                .addOption(OPTION_TABLE_SHORT, OPTION_TABLE_LONG, true, "Table name")
                .addOption(OPTION_CRON_EXPRESSION_SHORT, OPTION_CRON_EXPRESSION_LONG, true, "Cron expression to run extraction")
                .addOption(OPTION_HELP_SHORT, OPTION_HELP_LONG, false, "Print help message");
    }

    private static void printUsage(Options options) {
        HelpFormatter formatter = new HelpFormatter();
        formatter.printHelp(TrinoExtractor.class.getName(), options);
    }

    @DisallowConcurrentExecution
    public static class MetadataJob implements Job {
        private final Logger logger = LoggerFactory.getLogger(MetadataJob.class);

        public void execute(JobExecutionContext context) throws JobExecutionException {
            ExtractorContext extractorContext = TrinoExtractor.instance != null ? TrinoExtractor.instance.extractorContext : null;

            if (extractorContext != null) {
                logger.info("Executing metadata extraction at: {}", java.time.LocalTime.now());

                ExtractorService extractorService = new ExtractorService();

                try {
                    if (extractorService.execute(extractorContext)) {
                        TrinoExtractor.instance.exitCode = EXIT_CODE_SUCCESS;
                    }
                } catch (Exception e) {
                    logger.error("Error encountered: ", e);
                    throw new JobExecutionException(e);
                }
                logger.info("Completed executing metadata extraction at: {}", java.time.LocalTime.now());
            }
        }
    }
}
