# -*- coding: utf-8 -*-

#

# PgDiff - capture differences of database metadata

#

# Depends on PgBrowser

#

# Copyright (c) 2011  Radek Brich <radek.brich@devl.cz>

#

# Permission is hereby granted, free of charge, to any person obtaining a copy

# of this software and associated documentation files (the "Software"), to deal

# in the Software without restriction, including without limitation the rights

# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell

# copies of the Software, and to permit persons to whom the Software is

# furnished to do so, subject to the following conditions:

#

# The above copyright notice and this permission notice shall be included in

# all copies or substantial portions of the Software.

#

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR

# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,

# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE

# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER

# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,

# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN

# THE SOFTWARE.

from pycolib.ansicolor import *

import re

import difflib

class PgDiffError(Exception):

    pass

class DiffBase:

    COLORS = {

        '+' : BOLD | GREEN,

        '-' : BOLD | RED,

        '*' : BOLD | YELLOW,

    }

    COMMANDS = {

        '+' : 'CREATE',

        '-' : 'DROP',

        '*' : 'ALTER',

    }

    def __init__(self):

        self.changes = None

    def format(self):

        out = ['  ' * self.level]

        out.append(highlight(1, self.COLORS[self.change]))

        out.append(self.change)

        out += [' ', self.type, ' ', self.name, highlight(0)]

        if self.changes:

            out += [highlight(1, WHITE), ' (', self._formatchanges(), ')', highlight(0)]

        return ''.join(out)

    def _formatnotnull(self, notnull):

        if notnull:

            return 'NOT NULL'

        else:

            return None

    def _formatchanges(self):

        res = []

        for type, a, b in self.changes:

            if type == 'notnull':

                type = ''

                a = self._formatnotnull(a)

                b = self._formatnotnull(b)

            if a and b:

                s = ''.join(['Changed ', type, ' from ',

                    highlight(1,15), a, highlight(0), ' to ',

                    highlight(1,15), b, highlight(0), '.'])

            elif a and not b:

                l = ['Removed ']

                if type:

                    l += [type, ' ']

                l += [highlight(1,15), a, highlight(0), '.']

                s = ''.join(l)

            elif b and not a:

                l = ['Added ']

                if type:

                    l += [type, ' ']

                l += [highlight(1,15), b, highlight(0), '.']

                s = ''.join(l)

            res.append(s)

        return ' '.join(res)

    def format_patch(self):

        if self.change == '*' and self.type in ('schema', 'table'):

            return None

        return ['%s %s %s;' % (self.COMMANDS[self.change], self.type.upper(), self.name)]

class DiffSchema(DiffBase):

    def __init__(self, change, schema):

        DiffBase.__init__(self)

        self.level = 0

        self.type = 'schema'

        self.change = change

        self.schema = schema

        self.name = schema

class DiffTable(DiffBase):

    def __init__(self, change, schema, table):

        DiffBase.__init__(self)

        self.level = 1

        self.type = 'table'

        self.change = change

        self.schema = schema

        self.table = table

        self.name = table

class DiffArgument(DiffBase):

    def __init__(self, change, schema, function, argument):

        DiffBase.__init__(self)

        self.level = 2

        self.type = 'argument'

        self.change = change

        self.schema = schema

        self.function = function

        self.argument = argument

        self.name = argument

class DiffFunction(DiffBase):

    def __init__(self, change, schema, function, definition, show_body_diff=False):

        DiffBase.__init__(self)

        self.level = 1

        self.type = 'function'

        self.change = change

        self.schema = schema

        self.function = function

        #: New function definition

        self.definition = definition

        self.name = function

        self.show_body_diff = show_body_diff

    def _formatchanges(self):

        res = []

        for x in self.changes:

            type, a, b = x

            if type == 'source':

                if self.show_body_diff:

                    lines = ['Source differs:\n']

                    for line in difflib.unified_diff(a, b, lineterm=''):

                        if line[:3] in ('---', '+++'):

                            continue

                        color = {' ': WHITE, '-': YELLOW, '+': GREEN, '@': WHITE|BOLD}[line[0]]

                        lines.append(highlight(1, color) + line + highlight(0) + '\n')

                    res.append(''.join(lines))

                else:

                    res.append('Source differs.')

            else:

                res.append(''.join(['Changed ', type, ' from ',

                    highlight(1,15), a, highlight(0), ' to ',

                    highlight(1,15), b, highlight(0), '.']))

        return ' '.join(res)

    def format_patch(self):

        return [self.definition]

class DiffColumn(DiffBase):

    ALTER_COMMANDS = {

        '+' : 'ADD',

        '-' : 'DROP',

        '*' : 'ALTER',

    }

    def __init__(self, change, schema, table, column, columntype, columndefault, columnnotnull, changes=None):

        DiffBase.__init__(self)

        self.level = 2

        self.type = 'column'

        self.change = change

        self.schema = schema

        self.table = table

        self.column = column

        self.columntype = columntype

        self.columndefault = columndefault

        self.columnnotnull = columnnotnull

        self.name = column

        self.changes = changes

    def format_patch(self):

        alter_table = 'ALTER TABLE %s.%s %s COLUMN %s' % (

            self.schema,

            self.table,

            self.ALTER_COMMANDS[self.change],

            self.name,

        )

        out = []

        if self.change == '-':

            out.append('%s;' % alter_table);

        if self.change == '+':

            notnull = ''

            if self.columnnotnull:

                notnull = ' NOT NULL'

            default = ''

            if self.columndefault:

                default = ' DEFAULT %s' % self.columndefault

            out.append('%s %s%s%s;'

                % (alter_table, self.columntype, notnull, default));

        if self.change == '*':

            for type, a, b in self.changes:

                if type == 'type':

                    out.append('%s TYPE %s;' % (alter_table, b))

                if type == 'notnull':

                    if a and not b:

                        out.append('%s DROP NOT NULL;' % alter_table)

                    if not a and b:

                        out.append('%s SET NOT NULL;' % alter_table)

                if type == 'default':

                    if b:

                        out.append('%s SET DEFAULT %s;' % (alter_table, b))

                    else:

                        out.append('%s DROP DEFAULT;' % alter_table)

        return out

class DiffConstraint(DiffBase):

    def __init__(self, change, schema, table, constraint, definition, changes=None):

        DiffBase.__init__(self)

        self.level = 2

        self.type = 'constraint'

        self.change = change

        self.schema = schema

        self.table = table

        self.constraint = constraint

        self.name = constraint

        self.definition = definition

        self.changes = changes

    def format_patch(self):

        q_alter = 'ALTER TABLE %s.%s' % (self.schema, self.table)

        q_drop = '%s DROP CONSTRAINT %s;' % (q_alter, self.constraint)

        q_add = '%s ADD CONSTRAINT %s %s;' % (q_alter, self.constraint, self.definition)

        if self.change == '*':

            out = [q_drop, q_add]

        if self.change == '+':

            out = [q_add]

        if self.change == '-':

            out = [q_drop]

        return out

class DiffIndex(DiffBase):

    def __init__(self, change, schema, table, index, definition, changes=None):

        DiffBase.__init__(self)

        self.level = 2

        self.type = 'index'

        self.change = change

        self.schema = schema

        self.table = table

        self.index = index

        self.name = index

        self.definition = definition

        self.changes = changes

    def format_patch(self):

        q_drop = 'DROP INDEX %s;' % (self.index,)

        q_add = '%s;' % (self.definition,)

        if self.change == '*':

            out = [q_drop, q_add]

        if self.change == '+':

            out = [q_add]

        if self.change == '-':

            out = [q_drop]

        return out

class DiffType(DiffBase):

    def __init__(self, change, schema, name):

        DiffBase.__init__(self)

        self.level = 1

        self.type = 'type'

        self.change = change

        self.schema = schema

        self.name = name

class PgDiff:

    def __init__(self, srcbrowser=None, dstbrowser=None):

        self.allowcolor = False

        self.src = srcbrowser

        self.dst = dstbrowser

        self.include_schemas = set()  # if not empty, consider only these schemas for diff

        self.exclude_schemas = set()  # exclude these schemas from diff

        self.include_tables = set()

        self.exclude_tables = set()

        self.function_regex = re.compile(r"")

        self.function_body_diff = False

    def _test_schema(self, schema):

        if self.include_schemas and schema not in self.include_schemas:

            return False

        if schema in self.exclude_schemas:

            return False

        return True

    def _test_table(self, table):

        if self.include_tables and table not in self.include_tables:

            return False

        if table in self.exclude_tables:

            return False

        return True

    def _test_function(self, function):

        return bool(self.function_regex.match(function))

    def _diff_names(self, src, dst):

        for x in src:

            if x in dst:

                yield ('*', x)

            else:

                yield ('-', x)

        for x in dst:

            if x not in src:

                yield ('+', x)

    def _compare_columns(self, a, b):

        diff = []

        if a.type != b.type:

            diff.append(('type', a.type, b.type))

        if a.notnull != b.notnull:

            diff.append(('notnull', a.notnull, b.notnull))

        if a.default != b.default:

            diff.append(('default', a.default, b.default))

        return diff

    def _compare_constraints(self, a, b):

        diff = []

        if a.type != b.type:

            diff.append(('type', a.type, b.type))

        if a.definition != b.definition:

            diff.append(('definition', a.definition, b.definition))

        return diff

    def _compare_indexes(self, a, b):

        diff = []

        if a.definition != b.definition:

            diff.append(('definition', a.definition, b.definition))

        return diff

    def _compare_functions(self, a, b):

        diff = []

        if a.result != b.result:

            diff.append(('result', a.result, b.result))

        # function source may differ in newlines (\n vs \r\n)

        # split lines before comparison, so that these differencies are ignored

        a_source = a.source.splitlines()

        b_source = b.source.splitlines()

        if a_source != b_source:

            diff.append(('source', a_source, b_source))

        return diff

    def _compare_arguments(self, a, b):

        diff = []

        if a.type != b.type:

            diff.append(('type', a.type, b.type))

        if a.mode != b.mode:

            diff.append(('mode', a.mode, b.mode))

        if a.default != b.default:

            diff.append(('default', a.default, b.default))

        return diff

    def _compare_types(self, a, b):

        diff = []

        if a.type != b.type:

            diff.append(('type', a.type, b.type))

        if a.elements != b.elements:

            diff.append(('elements', repr(a.elements), repr(b.elements)))

        return diff

    def _diff_columns(self, schema, table, src_columns, dst_columns):

        for nd in self._diff_names(src_columns, dst_columns):

            if nd[1] in dst_columns:

                dst_type = dst_columns[nd[1]].type

                dst_default = dst_columns[nd[1]].default

                dst_notnull = dst_columns[nd[1]].notnull

            else:

                dst_type = None

                dst_default = None

                dst_notnull = None

            cdo = DiffColumn(change=nd[0], schema=schema, table=table, column=nd[1],

                columntype=dst_type, columndefault=dst_default, columnnotnull=dst_notnull)

            if nd[0] == '*':

                a = src_columns[nd[1]]

                b = dst_columns[nd[1]]

                cdo.changes = self._compare_columns(a, b)

                if cdo.changes:

                    yield cdo

            else:

                yield cdo

    def _diff_constraints(self, schema, table, src_constraints, dst_constraints):

        for nd in self._diff_names(src_constraints, dst_constraints):

            if nd[1] in dst_constraints:

                dst_definition = dst_constraints[nd[1]].definition

            else:

                dst_definition = None

            cdo = DiffConstraint(change=nd[0], schema=schema, table=table, constraint=nd[1],

                definition=dst_definition)

            if nd[0] == '*':

                a = src_constraints[nd[1]]

                b = dst_constraints[nd[1]]

                cdo.changes = self._compare_constraints(a, b)

                if cdo.changes:

                    yield cdo

            else:

                yield cdo

    def _diff_indexes(self, schema, table, src_indexes, dst_indexes):

        for nd in self._diff_names(src_indexes, dst_indexes):

            if nd[1] in dst_indexes:

                dst_definition = dst_indexes[nd[1]].definition

            else:

                dst_definition = None

            ido = DiffIndex(change=nd[0], schema=schema, table=table, index=nd[1],

                definition=dst_definition)

            if nd[0] == '*':

                a = src_indexes[nd[1]]

                b = dst_indexes[nd[1]]

                ido.changes = self._compare_indexes(a, b)

                if ido.changes:

                    yield ido

            else:

                yield ido

    def _diff_tables(self, schema, src_tables, dst_tables):

        for nd in self._diff_names(src_tables, dst_tables):

            if not self._test_table(nd[1]):

                continue

            tdo = DiffTable(change=nd[0], schema=schema, table=nd[1])

            if nd[0] == '*':

                # columns

                src_columns = src_tables[nd[1]].columns

                dst_columns = dst_tables[nd[1]].columns

                for cdo in self._diff_columns(schema, nd[1], src_columns, dst_columns):

                    if tdo:

                        yield tdo

                        tdo = None

                    yield cdo

                # constraints

                src_constraints = src_tables[nd[1]].constraints

                dst_constraints = dst_tables[nd[1]].constraints

                for cdo in self._diff_constraints(schema, nd[1], src_constraints, dst_constraints):

                    if tdo:

                        yield tdo

                        tdo = None

                    yield cdo

                # indexes

                src_indexes = src_tables[nd[1]].indexes

                dst_indexes = dst_tables[nd[1]].indexes

                for ido in self._diff_indexes(schema, nd[1], src_indexes, dst_indexes):

                    if tdo:

                        yield tdo

                        tdo = None

                    yield ido

            else:

                yield tdo

    def _diff_arguments(self, schema, function, src_args, dst_args):

        for nd in self._diff_names(src_args, dst_args):

            ado = DiffArgument(change=nd[0], schema=schema, function=function, argument=nd[1])

            if nd[0] == '*':

                a = src_args[nd[1]]

                b = dst_args[nd[1]]

                ado.changes = self._compare_arguments(a, b)

                if ado.changes:

                    yield ado

            else:

                yield ado

    def _diff_functions(self, schema, src_functions, dst_functions):

        for nd in self._diff_names(src_functions, dst_functions):

            if not self._test_function(nd[1]):

                continue

            if nd[1] in dst_functions:

                dst_definition = dst_functions[nd[1]].definition

            else:

                dst_definition = None

            fdo = DiffFunction(change=nd[0], schema=schema, function=nd[1],

                               definition=dst_definition,

                               show_body_diff=self.function_body_diff)

            if nd[0] == '*':

                # compare function body and result