/**
 * 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.
 */
/**
 * 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.hadoop.metadata.hive.hook;


import com.google.common.util.concurrent.ThreadFactoryBuilder;
import org.apache.hadoop.hive.conf.HiveConf;
import org.apache.hadoop.hive.ql.QueryPlan;
import org.apache.hadoop.hive.ql.exec.ExplainTask;
import org.apache.hadoop.hive.ql.exec.Task;
import org.apache.hadoop.hive.ql.hooks.Entity;
import org.apache.hadoop.hive.ql.hooks.ExecuteWithHookContext;
import org.apache.hadoop.hive.ql.hooks.HookContext;
import org.apache.hadoop.hive.ql.hooks.ReadEntity;
import org.apache.hadoop.hive.ql.hooks.WriteEntity;
import org.apache.hadoop.hive.ql.metadata.Table;
import org.apache.hadoop.hive.ql.plan.HiveOperation;
import org.apache.hadoop.metadata.hive.bridge.HiveMetaStoreBridge;
import org.apache.hadoop.metadata.hive.model.HiveDataTypes;
import org.apache.hadoop.metadata.typesystem.Referenceable;
import org.json.JSONObject;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.List;
import java.util.Set;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

/**
 * DgiHook sends lineage information to the DgiSever.
 */
public class HiveHook implements ExecuteWithHookContext {

    private static final Logger LOG = LoggerFactory.getLogger(HiveHook.class);

    // wait time determines how long we wait before we exit the jvm on
    // shutdown. Pending requests after that will not be sent.
    private static final int WAIT_TIME = 3;
    private static ExecutorService executor;

    private static final String MIN_THREADS = "hive.hook.dgi.minThreads";
    private static final String MAX_THREADS = "hive.hook.dgi.maxThreads";
    private static final String KEEP_ALIVE_TIME = "hive.hook.dgi.keepAliveTime";

    private static final int minThreadsDefault = 5;
    private static final int maxThreadsDefault = 5;
    private static final long keepAliveTimeDefault = 10;
    private static boolean typesRegistered = false;

    static {
        // anything shared should be initialized here and destroyed in the
        // shutdown hook The hook contract is weird in that it creates a
        // boatload of hooks.

        // initialize the async facility to process hook calls. We don't
        // want to do this inline since it adds plenty of overhead for the
        // query.
        HiveConf hiveConf = new HiveConf();
        int minThreads = hiveConf.getInt(MIN_THREADS, minThreadsDefault);
        int maxThreads = hiveConf.getInt(MAX_THREADS, maxThreadsDefault);
        long keepAliveTime = hiveConf.getLong(KEEP_ALIVE_TIME, keepAliveTimeDefault);

        executor = new ThreadPoolExecutor(minThreads, maxThreads, keepAliveTime, TimeUnit.MILLISECONDS,
                new LinkedBlockingQueue<Runnable>(),
                new ThreadFactoryBuilder().setDaemon(true).setNameFormat("DGI Logger %d").build());

        try {
            Runtime.getRuntime().addShutdownHook(new Thread() {
                        @Override
                        public void run() {
                            try {
                                executor.shutdown();
                                executor.awaitTermination(WAIT_TIME, TimeUnit.SECONDS);
                                executor = null;
                            } catch (InterruptedException ie) {
                                LOG.info("Interrupt received in shutdown.");
                            }
                            // shutdown client
                        }
                    });
        } catch (IllegalStateException is) {
            LOG.info("Attempting to send msg while shutdown in progress.");
        }

        LOG.info("Created DGI Hook");
    }

    class HiveEvent {
        public HiveConf conf;

        public Set<ReadEntity> inputs;
        public Set<WriteEntity> outputs;

        public String user;
        public HiveOperation operation;
        public QueryPlan queryPlan;
        public HookContext.HookType hookType;
        public JSONObject jsonPlan;
    }

    @Override
    public void run(final HookContext hookContext) throws Exception {
        if (executor == null) {
            LOG.info("No executor running. Bail.");
            return;
        }

        // clone to avoid concurrent access
        final HiveEvent event = new HiveEvent();
        final HiveConf conf = new HiveConf(hookContext.getConf());
        boolean debug = conf.get("hive.hook.dgi.synchronous", "false").equals("true");

        event.conf = conf;
        event.inputs = hookContext.getInputs();
        event.outputs = hookContext.getOutputs();

        event.user = hookContext.getUserName() == null ? hookContext.getUgi().getUserName() : hookContext.getUserName();
        event.operation = HiveOperation.valueOf(hookContext.getOperationName());
        event.queryPlan = hookContext.getQueryPlan();
        event.hookType = hookContext.getHookType();

        //todo throws NPE
//        event.jsonPlan = getQueryPlan(event);
        event.jsonPlan = new JSONObject();

        if (debug) {
            fireAndForget(event);
        } else {
            executor.submit(new Runnable() {
                        @Override
                        public void run() {
                            try {
                                fireAndForget(event);
                            } catch (Throwable e) {
                                LOG.info("DGI hook failed", e);
                            }
                        }
                    });
        }
    }

    private void fireAndForget(HiveEvent event) throws Exception {
        assert event.hookType == HookContext.HookType.POST_EXEC_HOOK : "Non-POST_EXEC_HOOK not supported!";

        LOG.info("Entered DGI hook for hook type {} operation {}", event.hookType, event.operation);
        HiveMetaStoreBridge dgiBridge = new HiveMetaStoreBridge(event.conf);

        if (!typesRegistered) {
            dgiBridge.registerHiveDataModel();
            typesRegistered = true;
        }

        switch (event.operation) {
        case CREATEDATABASE:
            handleCreateDB(dgiBridge, event);
            break;

        case CREATETABLE:
            handleCreateTable(dgiBridge, event);
            break;

        case CREATETABLE_AS_SELECT:
        case CREATEVIEW:
        case LOAD:
        case EXPORT:
        case IMPORT:
        case QUERY:
            registerProcess(dgiBridge, event);
            break;

        default:
        }
    }

    private void handleCreateTable(HiveMetaStoreBridge dgiBridge, HiveEvent event) throws Exception {
        for (WriteEntity entity : event.outputs) {
            if (entity.getType() == Entity.Type.TABLE) {

                Table table = entity.getTable();
                //TODO table.getDbName().toLowerCase() is required as hive stores in lowercase,
                // but table.getDbName() is not lowercase
                Referenceable dbReferenceable = dgiBridge.registerDatabase(table.getDbName().toLowerCase());
                dgiBridge.registerTable(dbReferenceable, table.getDbName(), table.getTableName());
            }
        }
    }

    private void handleCreateDB(HiveMetaStoreBridge dgiBridge, HiveEvent event) throws Exception {
        for (WriteEntity entity : event.outputs) {
            if (entity.getType() == Entity.Type.DATABASE) {
                dgiBridge.registerDatabase(entity.getDatabase().getName());
            }
        }
    }

    private void registerProcess(HiveMetaStoreBridge dgiBridge, HiveEvent event) throws Exception {
        Set<ReadEntity> inputs = event.inputs;
        Set<WriteEntity> outputs = event.outputs;

        //Even explain CTAS has operation name as CREATETABLE_AS_SELECT
        if (inputs.isEmpty() && outputs.isEmpty()) {
            LOG.info("Explain statement. Skipping...");
        }

        if (event.queryPlan == null) {
            LOG.info("Query plan is missing. Skipping...");
        }

        String queryId = event.queryPlan.getQueryId();
        String queryStr = event.queryPlan.getQueryStr();
        long queryStartTime = event.queryPlan.getQueryStartTime();

        LOG.debug("Registering CTAS query: {}", queryStr);
        Referenceable processReferenceable = new Referenceable(HiveDataTypes.HIVE_PROCESS.getName());
        processReferenceable.set("name", event.operation.getOperationName());
        processReferenceable.set("startTime", queryStartTime);
        processReferenceable.set("userName", event.user);
        List<Referenceable> source = new ArrayList<>();
        for (ReadEntity readEntity : inputs) {
            if (readEntity.getType() == Entity.Type.TABLE) {
                Table table = readEntity.getTable();
                String dbName = table.getDbName().toLowerCase();
                source.add(dgiBridge.registerTable(dbName, table.getTableName()));
            }
        }
        processReferenceable.set("inputTables", source);
        List<Referenceable> target = new ArrayList<>();
        for (WriteEntity writeEntity : outputs) {
            if (writeEntity.getType() == Entity.Type.TABLE || writeEntity.getType() == Entity.Type.PARTITION) {
                Table table = writeEntity.getTable();
                String dbName = table.getDbName().toLowerCase();
                target.add(dgiBridge.registerTable(dbName, table.getTableName()));
            }
            if (writeEntity.getType() == Entity.Type.PARTITION) {
                dgiBridge.registerPartition(writeEntity.getPartition());
            }
        }
        processReferenceable.set("outputTables", target);
        processReferenceable.set("queryText", queryStr);
        processReferenceable.set("queryId", queryId);
        processReferenceable.set("queryPlan", event.jsonPlan.toString());
        processReferenceable.set("endTime", System.currentTimeMillis());

        //TODO set
        processReferenceable.set("queryGraph", "queryGraph");
        dgiBridge.createInstance(processReferenceable);
    }


    private JSONObject getQueryPlan(HiveEvent event) throws Exception {
        ExplainTask explain = new ExplainTask();
        explain.initialize(event.conf, event.queryPlan, null);
        List<Task<?>> rootTasks = event.queryPlan.getRootTasks();
        return explain.getJSONPlan(null, null, rootTasks, event.queryPlan.getFetchTask(), true, false, false);
    }
}
