# Copyright (C) 2015, 2018, 2020, 2021 The Meme Factory, Inc.
# http://www.karlpinc.com/

# This file is part of PGWUI_Copy.
#
# This program is free software: you can redistribute it and/or
# modify it under the terms of the GNU Affero General Public License
# as published by the Free Software Foundation, either version 3 of
# the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public
# License along with this program.  If not, see
# <http://www.gnu.org/licenses/>.
#

import attr

import logging
import subprocess
import tempfile

import markupsafe

from wtforms import (
    BooleanField,
    StringField,
)

from pyramid.view import view_config

from pgwui_core.core import (
    UploadEngine,
    NoTransactionEngine,
    UserInitialPost,
    UserWTForm,
    CredsLoadedForm,
    SessionDBHandler,
    LogSQLCommand,
    SQLData,
    NoOpProcessor,
    ExecuteSQL,
    format_exception,
)

from pgwui_core.exceptions import AuthFailError

from pgwui_core.constants import (
    CHECKED,
    UNCHECKED,
)

from pgwui_common.view import errors_base_view

from pgwui_copy import utils

from pgwui_copy import exceptions as copy_ex

log = logging.getLogger(__name__)


@attr.s
class CopySchemaInitialPost(UserInitialPost):
    '''Get values from settings when there's not user-supplied values
    in POST
    '''
    schema = attr.ib(default=None)
    from_db = attr.ib(default=None)
    to_db = attr.ib(default=None)
    force = attr.ib(default=False)

    def build(self, settings={}):
        super().build(settings)
        copy_settings = settings['pgwui']['pgwui_copy']
        self.schema = copy_settings['default_schema']
        self.from_db = copy_settings['default_source_db']
        self.to_db = copy_settings['default_target_db']

        return self


class CopySchemaWTForm(UserWTForm):
    '''The wtform used to connect to the "gombemi" db to authenticate .'''
    # We don't actually use the labels, wanting the template to
    # look (and render) like html, but I'll define them anyway
    # just to keep my hand in.
    schema = StringField('Schema To Copy:')
    from_db = StringField('From Database:')
    to_db = StringField('To Database:')
    overwrite = BooleanField('Overwrite Destination Schema:')
    force = BooleanField('Allow Copy To Sensitive DBs:')


@attr.s
class CopySchemaForm(CredsLoadedForm):
    '''
    Acts like a dict, but with extra methods.

    Attributes:
      uh      The UploadHandler instance using the form
    '''
    def read(self):
        '''
        Read form data from the client
        '''

        # Read parent's data
        super().read()

        # Read our own data
        self['schema'] = self._form.schema.data.lower().strip()
        self['from_db'] = self._form.from_db.data.lower().strip()
        self['to_db'] = self._form.to_db.data.lower().strip()
        self['overwrite'] = self._form.overwrite.data
        self['force'] = self._form.force.data

    def write(self, result, errors):
        '''
        Produces the dict pyramid will use to render the form.
        '''
        response = super().write(result, errors)

        response['schema'] = self['schema']
        response['from_db'] = self['from_db']
        response['to_db'] = self['to_db']

        if self['force']:
            force_checked = CHECKED
        else:
            force_checked = UNCHECKED
        response['force'] = force_checked

        if self['overwrite']:
            overwrite_checked = CHECKED
        else:
            overwrite_checked = UNCHECKED
        response['overwrite'] = overwrite_checked

        return response


# Utility functions

def schema_exists(cur, schema):
    '''Does the schema exist?'''
    cur.execute('SELECT 1 FROM pg_namespace'
                '  WHERE nspname = %s',
                (schema,))
    return cur.fetchone() is not None


@attr.s
class CopySchemaHandler(SessionDBHandler):
    '''
    Abstract class to copy a schema from one db to another.
    '''
    def make_form(self):
        return CopySchemaForm().build(
            self, fc=CopySchemaWTForm,
            ip=CopySchemaInitialPost())

    def write(self, result, errors):
        '''
        Setup dict to render resulting html form

        Returns:
          Dict pyramid will use to render the resulting form
          Reserved keys:
            errors      A list of UploadError exceptions.
            db_changed  Boolean. Whether the db was changed.
        '''
        response = super().write(result, errors)

        response['db_changed'] = (not response['errors']
                                  and self.uf['action'] != '')

        return response

    def factory(self, ue):
        '''Make a db loader function from an UploadEngine.

        Input:

        Side Effects:
        Yes, lots.
        '''
        return ExecuteSQL(ue, self)

    def render(self, errors, result):
        '''Instead of rendering, just our results so we can
        decide what to do next.

        Input:
          errors    List of Error instances
          result    Db connection result dict
        '''
        response = super().render(errors, result)
        return (response, errors)

    def to_db_detail(self):
        '''Return a string for the log detailing an SQL execution attempt.'''
        return ('In database {db}: By user ({user})'
                .format(db=self.uf['to_db'], user=self.uf['user']))

    def log_sql_success(self, msg):
        '''Write a success message to the log when executing sql.'''
        log.warning('{0}: {1}'
                    .format(msg, self.to_db_detail()))

    def log_sql_failure(self, msg, ex):
        '''Write a failture message to the log when executing sql.'''
        log.info('{0}: {1}{2}'
                 .format(msg, self.to_db_detail(), format_exception(ex)))


@attr.s
class FromSchemaHandler(CopySchemaHandler):
    '''Handler to check the source schema and do initial validation.

    Attributes:
       cur
    '''
    settings = attr.ib(default=None)

    def init(self):
        super().init()
        self.settings = self.request.registry.settings

    def factory(self, ue):
        '''Make a db loader function from an UploadEngine.
        Having the UploadEngine we can test the db content.

        Input:

        Side Effects:
        Yes, lots.
        '''
        self.cur = ue.cur

        uf = self.uf
        cur = self.cur

        # Is the logged in user a superuser?
        user = uf['user']
        cur.execute('SELECT 1 FROM pg_roles'
                    '  WHERE rolname = %s AND rolsuper',
                    (user,))
        if cur.fetchone() is None:
            ex = copy_ex.NonAdminUserError(
                'Must be a superuser',
                ('The supplied user ({0}) is not an administrator'
                 .format(user)))
            self.log_sql_failure('User not a superuser', ex)
            raise ex

        schema = uf['schema']
        if not schema_exists(cur, schema):
            ex = copy_ex.InvalidSchemaError(
                'Source schema does not exist',
                ('The supplied source schema ({0}) does not exist'
                 .format(schema)))
            self.log_sql_failure('Source schema does not exist', ex)
            raise ex

        return super().factory(ue)

    def get_data(self):
        '''We do all the work validating input so need no data.'''
        self.data = ()

    def val_input(self):
        '''
        Validate input needed beyond that required to connect to the db.

        Returns:
          A list of Error instances
        '''
        uf = self.uf
        errors = []

        schema = uf['schema']

        # Do we have the essential basic data?
        if schema == '':
            ex = copy_ex.NoSchemaError('No schema supplied')
            self.log_sql_failure('No schema supplied', ex)
            errors.append(ex)

        sensitive_dbs = self.settings['pgwui']['pgwui_copy']['sensitive_dbs']
        if uf['to_db'] == '':
            ex = copy_ex.NoToDBError('No to db supplied')
            self.log_sql_failure('No to db supplied', ex)
            errors.append(ex)
        elif uf['to_db'].lower() in sensitive_dbs and not uf['force']:
            ex = copy_ex.LiveDBTargetError(
                'Cannot write into a sensitive database',
                'The "Allow ..." box must be checked to copy into '
                'the "{0}" database.'
                .format(uf['to_db']))
            self.log_sql_failure(
                'Aborted: The "Allow..." box must be checked to copy into '
                f'the "{uf["to_db"]}" database', ex)
            errors.append(ex)

        if uf['from_db'] == '':
            ex = copy_ex.NoFromDBError('No from db supplied')
            self.log_sql_failure('No from db supplied', ex)
            errors.append(ex)
        elif uf['from_db'] == uf['to_db']:
            ex = copy_ex.SameDBError(
                'The from db is the same as the to db',
                'Both the from db and the to db are ({0}).'
                .format(markupsafe.escape(uf['from_db'])))
            self.log_sql_failure('From db and to db are the same', ex)
            errors.append(ex)
        return errors


@attr.s
class DropSchemaHandler(CopySchemaHandler):
    '''
    Handler to drop the target schema from the target db.

    (Extra) Attributes:
      drop_attempted  True when the handler will attempt to drop the
                      schema.
    '''
    drop_attempted = attr.ib(default=False)

    def factory(self, ue):
        '''Make a db loader function from an UploadEngine.
        Having the UploadEngine we can test the db content.

        Input:

        Side Effects:
        Yes, lots.
        '''
        cur = ue.cur
        schema = self.uf['schema']
        if schema_exists(cur, schema):
            if self.uf['overwrite']:
                self.drop_attempted = True
                return super().factory(ue)
            else:
                log.info('Abort:'
                         ' Schema ({0}) already exists in the ({1}) db'
                         .format(schema, self.uf['to_db']))
                raise copy_ex.SchemaExistsError(
                    'The target schema exists',
                    'The schema ({0}) exists in the {1} db.'
                    .format(markupsafe.escape(schema), self.uf['to_db']))
        else:
            log.info('Schema ({0}) does not exist in the ({1}) db'
                     .format(schema, self.uf['to_db']))
            return NoOpProcessor(ue, self)

    def log_dropschema_success(self):
        '''Write a success message to the log when dropping the schema.'''
        self.log_sql_success('Success dropping schema ({0})'
                             .format(self.uf['schema']))

    def log_dropschema_failure(self, ex):
        '''Write a failure message to the log when dropping the schema.'''
        self.log_sql_failure('Failure dropping schema ({0})'
                             .format(self.uf['schema']),
                             ex)

    def log_vacuumfull_success(self):
        '''Write a success message to the log when vacuum full succeeds.'''
        self.log_sql_success('Successful VACUUM FULL')

    def log_vacuumfull_failure(self, ex):
        '''Write a failure message to the log when vacuum full fails.'''
        self.log_sql_failure('Failed VACUUM FULL', ex)

    def get_data(self):
        '''
        Build and stash the SQL to be executed.

        Returns:
          List of SQCommand instances
        '''
        uf = self.uf
        sql = []

        schema = uf['schema']
        to_db = uf['to_db']

        # Drop the schema
        sql.append(LogSQLCommand(
            'DROP SCHEMA {0} CASCADE'.format(schema),
            (),
            lambda ex: copy_ex.DropSchemaError(ex, schema, to_db),
            log_success=self.log_dropschema_success,
            log_failure=self.log_dropschema_failure))

        # Vacuum the target db
        sql.append(LogSQLCommand(
            'VACUUM FULL',
            (),
            lambda ex: copy_ex.VacuumFullError(ex, to_db),
            log_success=self.log_vacuumfull_success,
            log_failure=self.log_vacuumfull_failure))

        self.data = SQLData(sql)


class VacuumHandler(CopySchemaHandler):
    '''
    Handler to add a user schema to the test db.
    '''
    def log_vacuumanalyze_success(self):
        self.log_sql_success('Successful VACUUM ANALYZE')

    def log_vacuumanalyze_failure(self, ex):
        self.log_sql_failure('Failed VACUUM ANALYZE', ex)

    def get_data(self):
        '''
        Build and stash the SQL to be executed.

        Returns:
          List of SQCommand instances
        '''

        uf = self.uf
        sql = []

        to_db = uf['to_db']

        # Vacuum the target db
        sql.append(LogSQLCommand(
            'VACUUM ANALYZE',
            (),
            lambda ex: copy_ex.VacuumAnalyzeError(ex, to_db),
            log_success=self.log_vacuumanalyze_success,
            log_failure=self.log_vacuumanalyze_failure))

        self.data = SQLData(sql)


class CheckFromSchemaEngine(UploadEngine):
    '''Customize credential failure error message.'''
    def authfailerror_factory(self):
        orig = super().authfailerror_factory()
        return AuthFailError(
            ('Unable to login to the ({0}) db'
             .format(self.uh.uf['from_db'])),
            orig.descr)

    def read_db(self):
        '''The handler has two db's, save the right one.'''
        self.db = self.uh.uf['from_db']


class ToDBEngine(NoTransactionEngine):
    '''Use the to_db to login.'''
    def read_db(self):
        self.db = self.uh.uf['to_db']


def copy_schema(uf, settings):
    '''Copy the desired schema using pg_dump/pg_restore.

    Input:
      uf        The upload form, holding some of the info we need
                to do the copy.
      settings  The configuration settings, holding some of the info we need
                to do the copy.

    Returns:
      Exception, or None.
    '''
    host = settings['pg_host']
    port = settings['pg_port']
    bin = settings['pgwui_copy']['bin']
    user = uf['user']
    password = uf['password']
    schema = uf['schema']
    from_db = uf['from_db']
    to_db = uf['to_db']

    # Create the service file.
    sfile = tempfile.NamedTemporaryFile(mode='w',
                                        prefix='copy_schema_service_file')
    servicefile = sfile.name
    # No point in writing the dbname in [to_db] since (with 9.2 anyway)
    # pg_restore does not use it.
    sfile.writelines(
        '[from_db]\n'
        'host={host}\n'
        'port={port}\n'
        'dbname={from_db}\n'
        'user={user}\n'
        'password={password}\n'
        '\n'
        '[to_db]\n'
        'host={host}\n'
        'port={port}\n'
        'user={user}\n'
        'password={password}\n'
        .format(from_db=from_db,
                user=user,
                password=password,
                host=host,
                port=port))
    sfile.flush()

    # Run the necessary commands.
    log_info = ('User ({0}): From db ({1}): To db ({2}): Schema ({3})'
                .format(user, from_db, to_db, schema))
    try:
        pg_dump = subprocess.Popen(
            ('{0}/pg_dump'.format(bin),
             '--format', 'c', '--compress', '0',
             '--no-password',
             '--schema', schema,
             '--blobs'),
            bufsize=-1, close_fds=True,
            env={'PGSERVICEFILE': servicefile,
                 'PGSERVICE': 'from_db'},
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE)
    except OSError as ex:
        log.warn('Failed to execute pg_dump command: {0}: OS error {1}'
                 .format(log_info, ex.strerror))
        raise copy_ex.PGDumpOSError('Failed to execute pg_dump command', ex)

    try:
        pg_restore = subprocess.Popen(
            ('{0}/pg_restore'.format(bin),
             '--single-transaction',
             '--no-password',
             '--dbname', to_db),
            bufsize=-1, close_fds=True,
            env={'PGSERVICEFILE': servicefile,
                 'PGSERVICE': 'to_db'},
            stdin=pg_dump.stdout,
            stderr=subprocess.PIPE)
    except OSError as ex:
        log.warn('Failed to execute pg_restore command: {0}: OS error {1}'
                 .format(log_info, ex.strerror))
        raise copy_ex.PGRestoreOSError(
            'Failed to execute pg_restore command', ex)

    restore_errors = markupsafe.escape(
        pg_restore.communicate()[1].decode(errors='backslashreplace'))
    dump_errors = markupsafe.escape(
        pg_dump.communicate()[1].decode(errors='backslashreplace'))

    dump_retcode = pg_dump.returncode
    if dump_retcode:
        log.warn(f'pg_dump process failed: {log_info}:'
                 f' {utils.translate_retcode(dump_retcode)}:'
                 f' {dump_errors}')
        raise copy_ex.PGDumpError(
            ('Failure of pg_dump process to dump into '
             f'the "{from_db}" database'),
            dump_retcode, dump_errors)

    restore_retcode = pg_restore.returncode
    if restore_retcode:
        log.warn(f'pg_restore process failed: {log_info}:'
                 f' {utils.translate_retcode(restore_retcode)}:'
                 f' {restore_errors}')
        raise copy_ex.PGRestoreError(
            ('Failure of pg_restore process to restore into '
             f'the "{to_db}" database'),
            restore_retcode, restore_errors)

    log.warn('Successfully copied schema: {0}'.format(log_info))

    # Get rid of the service file.
    sfile.close()


def add_pgwui(response, settings):
    response['pgwui'] = settings['pgwui']
    return response


@view_config(route_name='pgwui_copy',
             renderer='pgwui_copy:templates/copy.mak')
@errors_base_view
def copy_schema_view(request):

    # Log that we started doing something.
    # We don't have a form yet, so get data right out of POST.
    post = request.POST
    form_submitted = 'action' in post and post['action'] != ''
    if form_submitted:
        # The form is submitted.
        log.debug('Processing form submission')

    # Check that the source schema exists and can be copied.
    # This is a normal engine and checks for csrf.
    from_handler = FromSchemaHandler(request)
    from_handler.init()
    response, errors = CheckFromSchemaEngine(from_handler).run()
    settings = request.registry.settings
    if errors:
        log.debug('Failed feasiblity test')
        return add_pgwui(response, settings)

    # Drop schema from the target db.
    drop_handler = DropSchemaHandler(request)
    drop_handler.init()
    response, errors = ToDBEngine(drop_handler).run()
    if response['db_changed']:
        # Copy schema to target db.
        try:
            copy_schema(from_handler.uf, settings['pgwui'])
        except copy_ex.CopyCommandError as ex:
            copyerrors = [ex]
            if drop_handler.drop_attempted:
                copyerrors.append(copy_ex.SchemaDroppedError(
                    'Target schema dropped but copy failed',
                    'The ({schema}) schema was dropped from the {db} db'
                    ' but the copy failed so now it no longer'
                    ' exists in the {db} db.'
                    .format(
                        schema=markupsafe.escape(
                            from_handler.uf['schema']),
                        db=markupsafe.escape(from_handler.uf['to_db']))))
            response['errors'].extend(copyerrors)
            response['db_changed'] = False
        else:
            schema = markupsafe.escape(from_handler.uf['schema'])
            to_db = markupsafe.escape(from_handler.uf['to_db'])
            from_db = markupsafe.escape(from_handler.uf['from_db'])
            user = markupsafe.escape(from_handler.uf['user'])
            log.info(f'Successfully copied schema:'
                     f' Schema ({schema}):'
                     f' From DB ({from_db}):'
                     f' To DB ({to_db}):'
                     f' By user {user}')

            # Vacuum analyze in the target db.
            vacuum_handler = VacuumHandler(request)
            vacuum_handler.init()
            vacuumout, vacuumerrors = ToDBEngine(vacuum_handler).run()
            # Merge results of modifying twice
            if vacuumerrors:
                vacuumerrors.append(copy_ex.SchemaCopiedError(
                    'VACUUM ANALYZE failed',
                    f'VACUUM ANALYZE failed in the {to_db} db,'
                    ' performance may suffer'))
                response['errors'].extend(vacuumerrors)
                del vacuumout['errors']
            else:
                log.info(f'Vacuum analyzed {to_db}')
            response.update(vacuumout)
    if form_submitted:
        log.debug('Finished processing request')
    return add_pgwui(response, settings)