"""
This is the MIT license: http://www.opensource.org/licenses/mit-license.php

Copyright (c) 2017 by Konstantin Lebedev.

Copyright 2022- 2023 Bytedance Ltd. and/or its affiliates

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 sqlalchemy import util as sa_util, types as sqltypes, exc
from sqlalchemy.sql import compiler, ClauseElement, expression
from sqlalchemy.sql.ddl import CreateColumn
from sqlalchemy.sql.elements import TextClause
from sqlalchemy.util import to_list


class ByteHouseDDLCompiler(compiler.DDLCompiler):
    def _get_default_string(self, default, name):
        sa_util.assert_arg_type(
            default, (sa_util.string_types[0], ClauseElement, TextClause), name
        )

        if isinstance(default, sa_util.string_types):
            return self.sql_compiler.render_literal_value(
                default, sqltypes.STRINGTYPE
            )
        else:
            return self.sql_compiler.process(
                default, literal_binds=True, include_table=False
            )

    def get_column_specification(self, column, **kw):
        colspec = (
            self.preparer.format_column(column)
            + " "
            + self.dialect.type_compiler.process(
                column.type, type_expression=column
            )
        )

        opts = column.dialect_options['bytehouse']

        if column.server_default is not None:
            colspec += " DEFAULT " + self._get_default_string(
                column.server_default.arg, 'server_default'
            )
        elif opts['materialized'] is not None:
            colspec += " MATERIALIZED " + self._get_default_string(
                opts['materialized'], 'bytehouse_materialized'
            )
        elif opts['alias'] is not None:
            colspec += " ALIAS " + self._get_default_string(
                opts['alias'], 'bytehouse_alias'
            )

        if column.comment is not None:
            literal = self.sql_compiler.render_literal_value(
                column.comment, sqltypes.String()
            )
            colspec += " COMMENT " + literal

        codec = opts['codec']
        if codec is not None:
            if isinstance(codec, (list, tuple)):
                codec = ', '.join(codec)
            colspec += " CODEC({0})".format(codec)

        if opts['after'] is not None:
            colspec += " AFTER " + self._get_default_string(
                opts['after'], 'bytehouse_after'
            )

        return colspec

    def visit_create_column(self, create, **kw):
        column = create.element
        nullable = column.nullable

        # All columns including synthetic PKs must be 'nullable'
        column.nullable = True

        rv = super(ByteHouseDDLCompiler, self).visit_create_column(
            create, **kw
        )
        column.nullable = nullable

        return rv

    def visit_primary_key_constraint(self, constraint, **kw):
        # Do not render PKs.
        return ''

    def _compile_param(self, expr, opt_list=False):
        compiler = self.sql_compiler

        # Do not render unnecessary brackets.
        if isinstance(expr, (list, tuple)) and len(expr) == 1 and opt_list:
            expr = expr[0]

        if isinstance(expr, (list, tuple)):
            return '(' + ', '.join(
                self._compile_param(el) for el in expr
            ) + ')'
        if not isinstance(expr, expression.ColumnClause):
            if not hasattr(expr, 'self_group'):
                # assuming base type (int, string, etc.)
                return str(expr)
            else:
                expr = expr.self_group()
        return compiler.process(
            expr, include_table=False, literal_binds=True
        )

    def visit_merge_tree(self, engine):
        param = engine.get_parameters()
        if param is None:
            param = '()'
        else:
            param = self._compile_param(to_list(param))

        text = '{0}{1}\n'.format(engine.name, param)
        if engine.partition_by:
            text += ' PARTITION BY {0}\n'.format(
                self._compile_param(
                    engine.partition_by.get_expressions_or_columns(),
                    opt_list=True
                )
            )
        if engine.order_by:
            text += ' ORDER BY {0}\n'.format(
                self._compile_param(
                    engine.order_by.get_expressions_or_columns(),
                    opt_list=True
                )
            )
        if engine.primary_key:
            text += ' PRIMARY KEY {0}\n'.format(
                self._compile_param(
                    engine.primary_key.get_expressions_or_columns(),
                    opt_list=True
                )
            )
        if engine.sample_by:
            text += ' SAMPLE BY {0}\n'.format(
                self._compile_param(
                    engine.sample_by.get_expressions_or_columns()[0]
                )
            )
        if engine.ttl:
            compile = self.sql_compiler.process
            text += ' TTL {0}\n'.format(
                ',\n     '.join(
                    compile(i, include_table=False, literal_binds=True)
                    for i in engine.ttl.get_expressions_or_columns()
                )
            )
        if engine.settings:
            text += ' SETTINGS ' + ', '.join(
                '{key}={value}'.format(
                    key=key,
                    value=value
                )
                for key, value in sorted(engine.settings.items())
            )
        return text

    def visit_engine(self, engine):
        engine_params = engine.get_parameters()
        text = engine.name
        if not engine_params:
            return text

        text += '('

        compiled_params = []
        for param in engine_params:
            if isinstance(param, tuple):
                compiled = (
                    '('
                    + ', '.join(self._compile_param(p) for p in param)
                    + ')'
                )
            else:
                compiled = self._compile_param(param)

            compiled_params.append(compiled)

        text += ', '.join(compiled_params)

        return text + ')'

    def create_table_suffix(self, table):
        cluster = table.dialect_options['bytehouse']['cluster']
        if cluster:
            return 'ON CLUSTER ' + self.preparer.quote(cluster)

        return ''

    def post_create_table(self, table):
        engine = getattr(table, 'engine', None)

        if not engine:
            raise exc.CompileError("No engine for table '%s'" % table.name)

        text = ' ENGINE = ' + self.process(engine)

        if table.comment is not None:
            literal = self.sql_compiler.render_literal_value(
                table.comment, sqltypes.String()
            )
            text += ' COMMENT ' + literal
        return text

    def visit_create_materialized_view(self, create):
        mv = create.element
        text = '\nCREATE MATERIALIZED VIEW '

        if create.if_not_exists:
            text += 'IF NOT EXISTS '

        text += self.preparer.format_table(mv)

        if mv.cluster:
            text += ' ON CLUSTER ' + self.preparer.quote(mv.cluster)

        if mv.to:
            text += ' TO ' + self.preparer.quote(mv.inner_table.name)

        else:
            text += ' ('

            separator = ''

            for column in mv.inner_table.columns:
                processed = self.process(CreateColumn(column))
                if processed is not None:
                    text += separator
                    separator = ', '
                    text += processed

            text += ')'

            text += self.post_create_table(mv.inner_table)

        text += ' '

        if mv.populate:
            text += 'POPULATE '

        text += 'AS ' + self.sql_compiler.process(
            mv.mv_selectable, literal_binds=True
        )

        return text

    def visit_drop_table(self, drop):
        text = '\nDROP TABLE '

        if getattr(drop, "if_exists", False):
            text += 'IF EXISTS '

        rv = text + self.preparer.format_table(drop.element)

        if getattr(drop, "on_cluster", False):
            rv += (
                ' ON CLUSTER ' +
                self.preparer.quote(drop.on_cluster)
            )

        return rv

    def visit_view(self, drop):
        text = '\nDROP VIEW '

        if getattr(drop, "if_exists", False):
            text += 'IF EXISTS '

        rv = text + self.preparer.format_table(drop.element)

        if getattr(drop, "on_cluster", False):
            rv += (
                ' ON CLUSTER ' +
                self.preparer.quote(drop.on_cluster)
            )

        return rv

    def visit_set_table_comment(self, create, **kw):
        return "ALTER TABLE %s MODIFY COMMENT %s" % (
            self.preparer.format_table(create.element),
            self.sql_compiler.render_literal_value(
                create.element.comment, sqltypes.String()
            )
        )

    def visit_drop_table_comment(self, create, **kw):
        return "ALTER TABLE %s MODIFY COMMENT ''" % (
            self.preparer.format_table(create.element)
        )
