##########################################################################
#
# pgAdmin 4 - PostgreSQL Tools
#
# Copyright (C) 2013 - 2021, The pgAdmin Development Team
# This software is released under the PostgreSQL Licence
#
##########################################################################

"""A blueprint module implementing the datagrid frame."""

import simplejson as json
import pickle
import random

from threading import Lock
from flask import Response, url_for, session, request, make_response
from werkzeug.useragents import UserAgent
from flask import current_app as app, render_template
from flask_babelex import gettext
from flask_security import login_required, current_user
from pgadmin.tools.sqleditor.command import ObjectRegistry, SQLFilter
from pgadmin.tools.sqleditor import check_transaction_status
from pgadmin.utils import PgAdminModule
from pgadmin.utils.ajax import make_json_response, bad_request, \
    internal_server_error, unauthorized

from config import PG_DEFAULT_DRIVER, ALLOW_SAVE_PASSWORD
from pgadmin.model import Server, User
from pgadmin.utils.driver import get_driver
from pgadmin.utils.exception import ConnectionLost, SSHTunnelConnectionLost
from pgadmin.utils.preferences import Preferences
from pgadmin.settings import get_setting
from pgadmin.browser.utils import underscore_unescape
from pgadmin.utils.exception import ObjectGone
from pgadmin.tools.sqleditor.utils.macros import get_user_macros
from pgadmin.utils.constants import MIMETYPE_APP_JS, UNAUTH_REQ

MODULE_NAME = 'datagrid'

query_tool_close_session_lock = Lock()


class DataGridModule(PgAdminModule):
    """
    class DataGridModule(PgAdminModule)

        A module class for Edit Grid derived from PgAdminModule.
    """

    LABEL = "Data Grid"

    def get_own_menuitems(self):
        return {}

    def get_own_javascripts(self):
        return [{
            'name': 'pgadmin.datagrid',
            'path': url_for('datagrid.index') + "datagrid",
            'when': None
        }]

    def get_panels(self):
        return []

    def get_exposed_url_endpoints(self):
        """
        Returns:
            list: URL endpoints for backup module
        """
        return [
            'datagrid.initialize_datagrid',
            'datagrid.initialize_query_tool',
            'datagrid.initialize_query_tool_with_did',
            'datagrid.filter_validate',
            'datagrid.filter',
            'datagrid.panel',
            'datagrid.close',
            'datagrid.update_query_tool_connection'
        ]

    def on_logout(self, user):
        """
        This is a callback function when user logout from pgAdmin
        :param user:
        :return:
        """
        with query_tool_close_session_lock:
            if 'gridData' in session:
                for trans_id in session['gridData']:
                    close_query_tool_session(trans_id)

                # Delete all grid data from session variable
                del session['gridData']


blueprint = DataGridModule(MODULE_NAME, __name__, static_url_path='/static')


@blueprint.route("/")
@login_required
def index():
    return bad_request(
        errormsg=gettext('This URL cannot be requested directly.')
    )


@blueprint.route("/css/datagrid.css")
def datagrid_css():
    return make_response(
        render_template('datagrid/css/datagrid.css'),
        200, {'Content-Type': 'text/css'}
    )


@blueprint.route("/filter", endpoint='filter')
@login_required
def show_filter():
    return render_template(MODULE_NAME + '/filter.html')


@blueprint.route(
    '/initialize/datagrid/<int:trans_id>/<int:cmd_type>/<obj_type>/'
    '<int:sgid>/<int:sid>/<int:did>/<int:obj_id>',
    methods=["PUT", "POST"],
    endpoint="initialize_datagrid"
)
@login_required
def initialize_datagrid(trans_id, cmd_type, obj_type, sgid, sid, did, obj_id):
    """
    This method is responsible for creating an asynchronous connection.
    After creating the connection it will instantiate and initialize
    the object as per the object type. It will also create a unique
    transaction id and store the information into session variable.

    Args:
        cmd_type: Contains value for which menu item is clicked.
        obj_type: Contains type of selected object for which data grid to
        be render
        sgid: Server group Id
        sid: Server Id
        did: Database Id
        obj_id: Id of currently selected object
    """

    if request.data:
        filter_sql = json.loads(request.data, encoding='utf-8')
    else:
        filter_sql = request.args or request.form

    # Create asynchronous connection using random connection id.
    conn_id = str(random.randint(1, 9999999))
    try:
        manager = get_driver(PG_DEFAULT_DRIVER).connection_manager(sid)
        # default_conn is same connection which is created when user connect to
        # database from tree
        default_conn = manager.connection(did=did)
        conn = manager.connection(did=did, conn_id=conn_id,
                                  auto_reconnect=False,
                                  use_binary_placeholder=True,
                                  array_to_string=True)
    except (ConnectionLost, SSHTunnelConnectionLost):
        raise
    except Exception as e:
        app.logger.error(e)
        return internal_server_error(errormsg=str(e))

    status, msg = default_conn.connect()
    if not status:
        app.logger.error(msg)
        return internal_server_error(errormsg=str(msg))

    status, msg = conn.connect()
    if not status:
        app.logger.error(msg)
        return internal_server_error(errormsg=str(msg))
    try:
        # if object type is partition then it is nothing but a table.
        if obj_type == 'partition':
            obj_type = 'table'

        # Get the object as per the object type
        command_obj = ObjectRegistry.get_object(
            obj_type, conn_id=conn_id, sgid=sgid, sid=sid,
            did=did, obj_id=obj_id, cmd_type=cmd_type,
            sql_filter=filter_sql
        )
    except ObjectGone:
        raise
    except Exception as e:
        app.logger.error(e)
        return internal_server_error(errormsg=str(e))

    if 'gridData' not in session:
        sql_grid_data = dict()
    else:
        sql_grid_data = session['gridData']

    # Use pickle to store the command object which will be used later by the
    # sql grid module.
    sql_grid_data[str(trans_id)] = {
        # -1 specify the highest protocol version available
        'command_obj': pickle.dumps(command_obj, -1)
    }

    # Store the grid dictionary into the session variable
    session['gridData'] = sql_grid_data

    return make_json_response(
        data={
            'conn_id': conn_id
        }
    )


@blueprint.route(
    '/panel/<int:trans_id>',
    methods=["POST"],
    endpoint='panel'
)
def panel(trans_id):
    """
    This method calls index.html to render the data grid.

    Args:
        trans_id: unique transaction id
    """

    url_params = None
    if request.args:
        url_params = {k: v for k, v in request.args.items()}

    if request.form:
        url_params['title'] = request.form['title']
        if 'sql_filter' in request.form:
            url_params['sql_filter'] = request.form['sql_filter']
        if 'query_url' in request.form:
            url_params['query_url'] = request.form['query_url']

    # We need client OS information to render correct Keyboard shortcuts
    user_agent = UserAgent(request.headers.get('User-Agent'))

    """
    Animations and transitions are not automatically GPU accelerated and by
    default use browser's slow rendering engine. We need to set 'translate3d'
    value of '-webkit-transform' property in order to use GPU. After applying
    this property under linux, Webkit calculates wrong position of the
    elements so panel contents are not visible. To make it work, we need to
    explicitly set '-webkit-transform' property to 'none' for .ajs-notifier,
    .ajs-message, .ajs-modal classes.

    This issue is only with linux runtime application and observed in Query
    tool and debugger. When we open 'Open File' dialog then whole Query tool
    panel content is not visible though it contains HTML element in back end.

    The port number should have already been set by the runtime if we're
    running in desktop mode.
    """
    is_linux_platform = False

    from sys import platform as _platform
    if "linux" in _platform:
        is_linux_platform = True

    # Fetch the server details
    bgcolor = None
    fgcolor = None

    s = Server.query.filter_by(id=url_params['sid']).first()
    if s and s.bgcolor:
        # If background is set to white means we do not have to change
        # the title background else change it as per user specified
        # background
        if s.bgcolor != '#ffffff':
            bgcolor = s.bgcolor
        fgcolor = s.fgcolor or 'black'

    layout = get_setting('SQLEditor/Layout')

    macros = get_user_macros()

    return render_template(
        "datagrid/index.html",
        _=gettext,
        uniqueId=trans_id,
        is_desktop_mode=app.PGADMIN_RUNTIME,
        is_linux=is_linux_platform,
        title=underscore_unescape(url_params['title']),
        url_params=json.dumps(url_params),
        client_platform=user_agent.platform,
        bgcolor=bgcolor,
        fgcolor=fgcolor,
        layout=layout,
        macros=macros
    )


@blueprint.route(
    '/initialize/query_tool/<int:trans_id>/<int:sgid>/<int:sid>/<int:did>',
    methods=["POST"], endpoint='initialize_query_tool_with_did'
)
@blueprint.route(
    '/initialize/query_tool/<int:trans_id>/<int:sgid>/<int:sid>',
    methods=["POST"], endpoint='initialize_query_tool'
)
@login_required
def initialize_query_tool(trans_id, sgid, sid, did=None):
    """
    This method is responsible for instantiating and initializing
    the query tool object. It will also create a unique
    transaction id and store the information into session variable.

    Args:
        sgid: Server group Id
        sid: Server Id
        did: Database Id
    """
    connect = True
    # Read the data if present. Skipping read may cause connection
    # reset error if data is sent from the client
    if request.data:
        _ = request.data

    req_args = request.args
    if ('recreate' in req_args and
            req_args['recreate'] == '1'):
        connect = False

    is_error, errmsg, conn_id, version = _init_query_tool(trans_id, connect,
                                                          sgid, sid, did)
    if is_error:
        return errmsg

    return make_json_response(
        data={
            'connId': str(conn_id),
            'serverVersion': version,
        }
    )


def _connect(conn, **kwargs):
    """
    Connect the database.
    :param conn: Connection instance.
    :param kwargs: user, role and password data from user.
    :return:
    """
    user = None
    role = None
    password = None
    is_ask_password = False
    if 'user' in kwargs and 'role' in kwargs:
        user = kwargs['user']
        role = kwargs['role'] if kwargs['role'] else None
        password = kwargs['password'] if kwargs['password'] else None
        is_ask_password = True
    if user:
        status, msg = conn.connect(user=user, role=role,
                                   password=password)
    else:
        status, msg = conn.connect()

    return status, msg, is_ask_password, user, role, password


def _init_query_tool(trans_id, connect, sgid, sid, did, **kwargs):
    # Create asynchronous connection using random connection id.
    conn_id = str(random.randint(1, 9999999))

    manager = get_driver(PG_DEFAULT_DRIVER).connection_manager(sid)

    if did is None:
        did = manager.did
    try:
        command_obj = ObjectRegistry.get_object(
            'query_tool', conn_id=conn_id, sgid=sgid, sid=sid, did=did
        )
    except Exception as e:
        app.logger.error(e)
        return True, internal_server_error(errormsg=str(e)), '', ''

    try:
        conn = manager.connection(did=did, conn_id=conn_id,
                                  auto_reconnect=False,
                                  use_binary_placeholder=True,
                                  array_to_string=True)

        if connect:
            status, msg, is_ask_password, user, role, password = _connect(
                conn, **kwargs)
            if not status:
                app.logger.error(msg)
                if is_ask_password:
                    server = Server.query.filter_by(id=sid).first()
                    return True, make_json_response(
                        success=0,
                        status=428,
                        result=render_template(
                            'servers/password.html',
                            server_label=server.name,
                            username=user,
                            errmsg=msg,
                            _=gettext,
                            allow_save_password=True if
                            ALLOW_SAVE_PASSWORD and
                            session['allow_save_password'] else False,
                        )
                    ), '', ''
                else:
                    return True, internal_server_error(
                        errormsg=str(msg)), '', ''
    except (ConnectionLost, SSHTunnelConnectionLost) as e:
        app.logger.error(e)
        raise
    except Exception as e:
        app.logger.error(e)
        return True, internal_server_error(errormsg=str(e)), '', ''

    if 'gridData' not in session:
        sql_grid_data = dict()
    else:
        sql_grid_data = session['gridData']

    # Set the value of auto commit and auto rollback specified in Preferences
    pref = Preferences.module('sqleditor')
    command_obj.set_auto_commit(pref.preference('auto_commit').get())
    command_obj.set_auto_rollback(pref.preference('auto_rollback').get())

    # Use pickle to store the command object which will be used
    # later by the sql grid module.
    sql_grid_data[str(trans_id)] = {
        # -1 specify the highest protocol version available
        'command_obj': pickle.dumps(command_obj, -1)
    }

    # Store the grid dictionary into the session variable
    session['gridData'] = sql_grid_data

    return False, '', conn_id, manager.version


@blueprint.route(
    '/initialize/query_tool/update_connection/<int:trans_id>/'
    '<int:sgid>/<int:sid>/<int:did>',
    methods=["POST"], endpoint='update_query_tool_connection'
)
def update_query_tool_connection(trans_id, sgid, sid, did):
    # Remove transaction Id.
    with query_tool_close_session_lock:
        data = json.loads(request.data, encoding='utf-8')

        if 'gridData' not in session:
            return make_json_response(data={'status': True})

        grid_data = session['gridData']

        # Return from the function if transaction id not found
        if str(trans_id) not in grid_data:
            return make_json_response(data={'status': True})

        connect = True

        req_args = request.args
        if ('recreate' in req_args and
                req_args['recreate'] == '1'):
            connect = False

        new_trans_id = str(random.randint(1, 9999999))
        kwargs = {
            'user': data['user'],
            'role': data['role'] if 'role' in data else None,
            'password': data['password'] if 'password' in data else None
        }

        is_error, errmsg, conn_id, version = _init_query_tool(
            new_trans_id, connect, sgid, sid, did, **kwargs)

        if is_error:
            return errmsg
        else:
            try:
                # Check the transaction and connection status
                status, error_msg, conn, trans_obj, session_obj = \
                    check_transaction_status(trans_id)

                status, error_msg, new_conn, new_trans_obj, new_session_obj = \
                    check_transaction_status(new_trans_id)

                new_session_obj['primary_keys'] = session_obj[
                    'primary_keys'] if 'primary_keys' in session_obj else None
                new_session_obj['columns_info'] = session_obj[
                    'columns_info'] if 'columns_info' in session_obj else None
                new_session_obj['client_primary_key'] = session_obj[
                    'client_primary_key'] if 'client_primary_key'\
                                             in session_obj else None

                close_query_tool_session(trans_id)
                # Remove the information of unique transaction id from the
                # session variable.
                grid_data.pop(str(trans_id), None)
                session['gridData'] = grid_data
            except Exception as e:
                app.logger.error(e)

    return make_json_response(
        data={
            'connId': str(conn_id),
            'serverVersion': version,
            'tran_id': new_trans_id
        }
    )


@blueprint.route('/close/<int:trans_id>', methods=["DELETE"], endpoint='close')
def close(trans_id):
    """
    This method is used to close the asynchronous connection
    and remove the information of unique transaction id from
    the session variable.

    Args:
        trans_id: unique transaction id
    """
    with query_tool_close_session_lock:
        if 'gridData' not in session:
            return make_json_response(data={'status': True})

        grid_data = session['gridData']
        # Return from the function if transaction id not found
        if str(trans_id) not in grid_data:
            return make_json_response(data={'status': True})

        try:
            close_query_tool_session(trans_id)
            # Remove the information of unique transaction id from the
            # session variable.
            grid_data.pop(str(trans_id), None)
            session['gridData'] = grid_data
        except Exception as e:
            app.logger.error(e)
            return internal_server_error(errormsg=str(e))

    return make_json_response(data={'status': True})


@blueprint.route(
    '/filter/validate/<int:sid>/<int:did>/<int:obj_id>',
    methods=["PUT", "POST"], endpoint='filter_validate'
)
@login_required
def validate_filter(sid, did, obj_id):
    """
    This method is used to validate the sql filter.

    Args:
        sid: Server Id
        did: Database Id
        obj_id: Id of currently selected object
    """
    if request.data:
        filter_sql = json.loads(request.data, encoding='utf-8')
    else:
        filter_sql = request.args or request.form

    try:
        # Create object of SQLFilter class
        sql_filter_obj = SQLFilter(sid=sid, did=did, obj_id=obj_id)

        # Call validate_filter method to validate the SQL.
        status, res = sql_filter_obj.validate_filter(filter_sql)
    except ObjectGone:
        raise
    except Exception as e:
        app.logger.error(e)
        return internal_server_error(errormsg=str(e))

    return make_json_response(data={'status': status, 'result': res})


@blueprint.route("/datagrid.js")
@login_required
def script():
    """render the required javascript"""
    return Response(
        response=render_template("datagrid/js/datagrid.js", _=gettext),
        status=200,
        mimetype=MIMETYPE_APP_JS
    )


def close_query_tool_session(trans_id):
    """
    This function is used to cancel the transaction and release the connection.

    :param trans_id: Transaction id
    :return:
    """
    if 'gridData' in session and str(trans_id) in session['gridData']:
        cmd_obj_str = session['gridData'][str(trans_id)]['command_obj']
        # Use pickle.loads function to get the command object
        cmd_obj = pickle.loads(cmd_obj_str)

        # if connection id is None then no need to release the connection
        if cmd_obj.conn_id is not None:
            manager = get_driver(
                PG_DEFAULT_DRIVER).connection_manager(cmd_obj.sid)
            if manager is not None:
                conn = manager.connection(
                    did=cmd_obj.did, conn_id=cmd_obj.conn_id)

                # Release the connection
                if conn.connected():
                    conn.cancel_transaction(cmd_obj.conn_id, cmd_obj.did)
                    manager.release(did=cmd_obj.did, conn_id=cmd_obj.conn_id)
