# -*- coding: utf-8 -*-
#
# 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
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# 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.

from datetime import datetime
import json

from flask import Blueprint, request
from flask_admin import BaseView, expose
import pandas as pd

from airflow.hooks.hive_hooks import HiveMetastoreHook, HiveCliHook
from airflow.hooks.mysql_hook import MySqlHook
from airflow.hooks.presto_hook import PrestoHook
from airflow.plugins_manager import AirflowPlugin
from airflow.www import utils as wwwutils
from airflow.www.decorators import gzipped

METASTORE_CONN_ID = 'metastore_default'
METASTORE_MYSQL_CONN_ID = 'metastore_mysql'
PRESTO_CONN_ID = 'presto_default'
HIVE_CLI_CONN_ID = 'hive_default'
DEFAULT_DB = 'default'
DB_WHITELIST = None
DB_BLACKLIST = ['tmp']
TABLE_SELECTOR_LIMIT = 2000

# Keeping pandas from truncating long strings
pd.set_option('display.max_colwidth', -1)


# Creating a flask admin BaseView
class MetastoreBrowserView(BaseView, wwwutils.DataProfilingMixin):

    @expose('/')
    def index(self):
        sql = """
        SELECT
            a.name as db, db_location_uri as location,
            count(1) as object_count, a.desc as description
        FROM DBS a
        JOIN TBLS b ON a.DB_ID = b.DB_ID
        GROUP BY a.name, db_location_uri, a.desc
        """.format(**locals())
        h = MySqlHook(METASTORE_MYSQL_CONN_ID)
        df = h.get_pandas_df(sql)
        df.db = (
            '<a href="/admin/metastorebrowserview/db/?db=' +
            df.db + '">' + df.db + '</a>')
        table = df.to_html(
            classes="table table-striped table-bordered table-hover",
            index=False,
            escape=False,
            na_rep='',)
        return self.render(
            "metastore_browser/dbs.html", table=table)

    @expose('/table/')
    def table(self):
        table_name = request.args.get("table")
        m = HiveMetastoreHook(METASTORE_CONN_ID)
        table = m.get_table(table_name)
        return self.render(
            "metastore_browser/table.html",
            table=table, table_name=table_name, datetime=datetime, int=int)

    @expose('/db/')
    def db(self):
        db = request.args.get("db")
        m = HiveMetastoreHook(METASTORE_CONN_ID)
        tables = sorted(m.get_tables(db=db), key=lambda x: x.tableName)
        return self.render(
            "metastore_browser/db.html", tables=tables, db=db)

    @gzipped
    @expose('/partitions/')
    def partitions(self):
        schema, table = request.args.get("table").split('.')
        sql = """
        SELECT
            a.PART_NAME,
            a.CREATE_TIME,
            c.LOCATION,
            c.IS_COMPRESSED,
            c.INPUT_FORMAT,
            c.OUTPUT_FORMAT
        FROM PARTITIONS a
        JOIN TBLS b ON a.TBL_ID = b.TBL_ID
        JOIN DBS d ON b.DB_ID = d.DB_ID
        JOIN SDS c ON a.SD_ID = c.SD_ID
        WHERE
            b.TBL_NAME like '{table}' AND
            d.NAME like '{schema}'
        ORDER BY PART_NAME DESC
        """.format(**locals())
        h = MySqlHook(METASTORE_MYSQL_CONN_ID)
        df = h.get_pandas_df(sql)
        return df.to_html(
            classes="table table-striped table-bordered table-hover",
            index=False,
            na_rep='',)

    @gzipped
    @expose('/objects/')
    def objects(self):
        where_clause = ''
        if DB_WHITELIST:
            dbs = ",".join(["'" + db + "'" for db in DB_WHITELIST])
            where_clause = "AND b.name IN ({})".format(dbs)
        if DB_BLACKLIST:
            dbs = ",".join(["'" + db + "'" for db in DB_BLACKLIST])
            where_clause = "AND b.name NOT IN ({})".format(dbs)
        sql = """
        SELECT CONCAT(b.NAME, '.', a.TBL_NAME), TBL_TYPE
        FROM TBLS a
        JOIN DBS b ON a.DB_ID = b.DB_ID
        WHERE
            a.TBL_NAME NOT LIKE '%tmp%' AND
            a.TBL_NAME NOT LIKE '%temp%' AND
            b.NAME NOT LIKE '%tmp%' AND
            b.NAME NOT LIKE '%temp%'
        {where_clause}
        LIMIT {LIMIT};
        """.format(where_clause=where_clause, LIMIT=TABLE_SELECTOR_LIMIT)
        h = MySqlHook(METASTORE_MYSQL_CONN_ID)
        d = [
            {'id': row[0], 'text': row[0]}
            for row in h.get_records(sql)]
        return json.dumps(d)

    @gzipped
    @expose('/data/')
    def data(self):
        table = request.args.get("table")
        sql = "SELECT * FROM {table} LIMIT 1000;".format(table=table)
        h = PrestoHook(PRESTO_CONN_ID)
        df = h.get_pandas_df(sql)
        return df.to_html(
            classes="table table-striped table-bordered table-hover",
            index=False,
            na_rep='',)

    @expose('/ddl/')
    def ddl(self):
        table = request.args.get("table")
        sql = "SHOW CREATE TABLE {table};".format(table=table)
        h = HiveCliHook(HIVE_CLI_CONN_ID)
        return h.run_cli(sql)


v = MetastoreBrowserView(category="Plugins", name="Hive Metadata Browser")

# Creating a flask blueprint to intergrate the templates and static folder
bp = Blueprint(
    "metastore_browser", __name__,
    template_folder='templates',
    static_folder='static',
    static_url_path='/static/metastore_browser')


# Defining the plugin class
class MetastoreBrowserPlugin(AirflowPlugin):
    name = "metastore_browser"
    flask_blueprints = [bp]
    admin_views = [v]

[(44, 'pandas.set_option', 'pd.set_option', True, 'import pandas as pd\n'), (169, 'flask.Blueprint', 'Blueprint', False, 'from flask import Blueprint, request\n'), (50, 'flask_admin.expose', 'expose', False, 'from flask_admin import BaseView, expose\n'), (73, 'flask_admin.expose', 'expose', False, 'from flask_admin import BaseView, expose\n'), (82, 'flask_admin.expose', 'expose', False, 'from flask_admin import BaseView, expose\n'), (91, 'flask_admin.expose', 'expose', False, 'from flask_admin import BaseView, expose\n'), (119, 'flask_admin.expose', 'expose', False, 'from flask_admin import BaseView, expose\n'), (147, 'flask_admin.expose', 'expose', False, 'from flask_admin import BaseView, expose\n'), (158, 'flask_admin.expose', 'expose', False, 'from flask_admin import BaseView, expose\n'), (60, 'airflow.hooks.mysql_hook.MySqlHook', 'MySqlHook', False, 'from airflow.hooks.mysql_hook import MySqlHook\n'), (75, 'flask.request.args.get', 'request.args.get', False, 'from flask import Blueprint, request\n'), (76, 'airflow.hooks.hive_hooks.HiveMetastoreHook', 'HiveMetastoreHook', False, 'from airflow.hooks.hive_hooks import HiveMetastoreHook, HiveCliHook\n'), (84, 'flask.request.args.get', 'request.args.get', False, 'from flask import Blueprint, request\n'), (85, 'airflow.hooks.hive_hooks.HiveMetastoreHook', 'HiveMetastoreHook', False, 'from airflow.hooks.hive_hooks import HiveMetastoreHook, HiveCliHook\n'), (111, 'airflow.hooks.mysql_hook.MySqlHook', 'MySqlHook', False, 'from airflow.hooks.mysql_hook import MySqlHook\n'), (140, 'airflow.hooks.mysql_hook.MySqlHook', 'MySqlHook', False, 'from airflow.hooks.mysql_hook import MySqlHook\n'), (144, 'json.dumps', 'json.dumps', False, 'import json\n'), (149, 'flask.request.args.get', 'request.args.get', False, 'from flask import Blueprint, request\n'), (151, 'airflow.hooks.presto_hook.PrestoHook', 'PrestoHook', False, 'from airflow.hooks.presto_hook import PrestoHook\n'), (160, 'flask.request.args.get', 'request.args.get', False, 'from flask import Blueprint, request\n'), (162, 'airflow.hooks.hive_hooks.HiveCliHook', 'HiveCliHook', False, 'from airflow.hooks.hive_hooks import HiveMetastoreHook, HiveCliHook\n'), (93, 'flask.request.args.get', 'request.args.get', False, 'from flask import Blueprint, request\n')]