#! /usr/bin/python3
# Copyright 2021-present StarRocks, Inc. All rights reserved.
#
# Licensed 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
#
#     https:#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.

import decimal
from textwrap import dedent

from sqlalchemy import VARCHAR, Column, Integer, MetaData, Table, literal, schema, testing
from sqlalchemy.sql.sqltypes import Float
from sqlalchemy.testing import fixtures
from sqlalchemy.testing.assertions import AssertsCompiledSQL, eq_
from sqlalchemy.testing.suite import *


class CompileTest(fixtures.TestBase, AssertsCompiledSQL):

    __only_on__ = "starrocks"

    def test_create_table_with_properties(self):
        m = MetaData()
        tbl = Table(
            'atable', m, Column("id", Integer),
            starrocks_properties=(
                ("storage_medium", "SSD"),
                ("storage_cooldown_time", "2015-06-04 00:00:00"),
            ))
        self.assert_compile(
            schema.CreateTable(tbl),
            "CREATE TABLE atable (id INTEGER)PROPERTIES(\"storage_medium\"=\"SSD\",\"storage_cooldown_time\"=\"2015-06-04 00:00:00\")")

    def test_create_primary_key_table(self):
        m = MetaData()
        tbl = Table(
            'btable', m, Column("id", Integer, primary_key=True),
            starrocks_primary_key="id",
            starrocks_distributed_by="id"
            )
        self.assert_compile(
            schema.CreateTable(tbl),
            "CREATE TABLE btable (id BIGINT NOT NULL AUTO_INCREMENT)PRIMARY KEY(id) DISTRIBUTED BY HASH(id)")


# Float test fixes below for "Data type of first column cannot be FLOAT" error given by starrocks
class _LiteralRoundTripFixture2(object):
    supports_whereclause = True

    @testing.fixture
    def literal_round_trip(self, metadata, connection):
        """test literal rendering"""

        # for literal, we test the literal render in an INSERT
        # into a typed column.  we can then SELECT it back as its
        # official type; ideally we'd be able to use CAST here
        # but MySQL in particular can't CAST fully

        def run(type_, input_, output, filter_=None):
            t = Table("t", metadata, Column("a", Integer), Column("x", type_))
            t.create(connection)

            for value in input_:
                ins = (
                    t.insert()
                    .values(x=literal(value, type_))
                    .compile(
                        dialect=testing.db.dialect,
                        compile_kwargs=dict(literal_binds=True),
                    )
                )
                connection.execute(ins)

            if self.supports_whereclause:
                stmt = t.select().where(t.c.x == literal(value))
            else:
                stmt = t.select()

            stmt = stmt.compile(
                dialect=testing.db.dialect,
                compile_kwargs=dict(literal_binds=True),
            )
            for row in connection.execute(stmt):
                value = row[0]
                if filter_ is not None:
                    value = filter_(value)
                assert value in output

        return run


class FloatTest(_LiteralRoundTripFixture2, fixtures.TestBase):
    __backend__ = True

    @testing.fixture
    def do_numeric_test(self, metadata, connection):
        @testing.emits_warning(
            r".*does \*not\* support Decimal objects natively"
        )
        def run(type_, input_, output, filter_=None, check_scale=False):
            t = Table("t", metadata, Column("a", Integer), Column("x", type_))
            t.create(connection)
            connection.execute(t.insert(), [{"x": x} for x in input_])

            result = {row[1] for row in connection.execute(t.select())}
            output = set(output)
            if filter_:
                result = set(filter_(x) for x in result)
                output = set(filter_(x) for x in output)
            eq_(result, output)
            if check_scale:
                eq_([str(x) for x in result], [str(x) for x in output])

        return run

    @testing.requires.floats_to_four_decimals
    def test_float_as_decimal(self, do_numeric_test):
        do_numeric_test(
            Float(precision=8, asdecimal=True),
            [15.7563, decimal.Decimal("15.7563"), None],
            [decimal.Decimal("15.7563"), None],
            filter_=lambda n: n is not None and round(n, 4) or None,
        )

    def test_float_as_float(self, do_numeric_test):
        do_numeric_test(
            Float(precision=8),
            [15.7563, decimal.Decimal("15.7563")],
            [15.7563],
            filter_=lambda n: n is not None and round(n, 5) or None,
        )

    @testing.requires.precision_generic_float_type
    def test_float_custom_scale(self, do_numeric_test):
        do_numeric_test(
            Float(None, decimal_return_scale=6, asdecimal=True),
            [15.756382, decimal.Decimal("15.756382")],
            [decimal.Decimal("15.756382")],
            check_scale=True,
        )

    def test_render_literal_float(self, literal_round_trip):
        literal_round_trip(
            Float(4),
            [15.7563, decimal.Decimal("15.7563")],
            [15.7563],
            filter_=lambda n: n is not None and round(n, 5) or None,
        )


class StarRocksReflectionTest(fixtures.TestBase):
    __only_on__ = "starrocks"

    # @testing.combinations(True, False, argnames="primary_key")
    def test_default_reflection(self, connection, metadata):
        tn = "test_starrocks_reflection"
        Table(
            tn,
            metadata,
            Column("id", Integer),
            Column("data", VARCHAR(10)),
        )
        metadata.create_all(connection)

        reflected_meta = MetaData()
        reflected_meta.reflect(connection, only=[tn])
        reflected_table = reflected_meta.tables[tn]

        eq_(reflected_table.name, tn)
        eq_(reflected_table.schema, None)
        eq_(len(reflected_table.columns), 2)
        eq_(reflected_table.comment, "")
        eq_(set(reflected_table.columns.keys()), {"id", "data"})
        assert isinstance(reflected_table.columns["id"].type, Integer)
        assert isinstance(reflected_table.columns["data"].type, VARCHAR)
        eq_(reflected_table.columns["data"].type.length, 10)
        for c in ["id", "data"]:
            eq_(reflected_table.columns[c].comment, "")
            eq_(reflected_table.columns[c].nullable, True)
            eq_(reflected_table.columns[c].default, None)
            eq_(reflected_table.columns[c].autoincrement, None)
        eq_(reflected_table.dialect_options["starrocks"]["comment"], "")
        eq_(reflected_table.dialect_options["starrocks"]["distribution"], "RANDOM")
        eq_(reflected_table.dialect_options["starrocks"]["engine"], "OLAP")
        eq_(reflected_table.dialect_options["starrocks"]["key_desc"], "DUPLICATE KEY(`id`, `data`)")
        eq_(reflected_table.dialect_options["starrocks"]["order_by"], "`id`, `data`")
        eq_(reflected_table.dialect_options["starrocks"]["partition_by"], "")
        assert ("compression", "LZ4") in reflected_table.dialect_options["starrocks"]["properties"]

    def test_comment_reflection(self, connection, metadata):
        tn = "test_starrocks_reflection_comment"

        connection.execute(text(f"DROP TABLE IF EXISTS {tn}"))
        create_table_sql = dedent(f"""
            CREATE TABLE {tn} (
                id INT COMMENT 'This is id column',
                data VARCHAR(10) COMMENT 'This is data column'
            )
            COMMENT 'This is a test table'
        """).strip()
        connection.execute(text(create_table_sql))

        reflected_meta = MetaData()
        reflected_meta.reflect(connection, only=[tn])
        reflected_table = reflected_meta.tables[tn]

        eq_(reflected_table.comment, "This is a test table")
        eq_(reflected_table.columns["id"].comment, "This is id column")
        eq_(reflected_table.columns["data"].comment, "This is data column")

    def test_view_comment_reflection(self, connection, metadata):
        vn = "test_starrocks_reflection_view_comment"

        connection.execute(text(f"DROP VIEW IF EXISTS {vn}"))
        create_table_sql = dedent(f"""
            CREATE VIEW {vn} (
                id COMMENT 'This is id column',
                data COMMENT 'This is data column'
            )
            COMMENT 'This is a test view'
            AS SELECT 1 AS id, 'data' AS data
        """).strip()
        connection.execute(text(create_table_sql))

        reflected_meta = MetaData()
        reflected_meta.reflect(connection, only=[vn], views=True)
        reflected_table = reflected_meta.tables[vn]

        eq_(reflected_table.comment, "This is a test view")
        eq_(reflected_table.columns["id"].comment, "This is id column")
        eq_(reflected_table.columns["data"].comment, "This is data column")

    @testing.combinations(["id"], ["id", "data"], argnames="key_columns")
    @testing.combinations("PRIMARY", "DUPLICATE", "UNIQUE", argnames="key_type")
    def test_key_and_distribution_reflection(self, connection, metadata, key_type, key_columns):
        tn = "test_starrocks_reflection_key"
        connection.execute(text(f"DROP TABLE IF EXISTS {tn}"))
        create_table_sql = dedent(f"""
            CREATE TABLE {tn} (
                id INT,
                data VARCHAR(10),
                something_else DATETIME
            )
            {key_type} KEY ({', '.join(key_columns)})
            DISTRIBUTED BY HASH(id)
        """).strip()
        connection.execute(text(create_table_sql))

        reflected_meta = MetaData()
        reflected_meta.reflect(connection, only=[tn])
        reflected_table = reflected_meta.tables[tn]

        if key_type == "PRIMARY":
            for c in key_columns:
                eq_(reflected_table.columns[c].nullable, False)
        eq_(reflected_table.dialect_options["starrocks"]["key_desc"], f"{key_type} KEY({', '.join([f'`{c}`' for c in key_columns])})")
        eq_(reflected_table.dialect_options["starrocks"]["distribution"], "HASH(`id`)")

    @testing.combinations(["id"], ["id", "data"], argnames="partition_by")
    def test_partition_by_columns_reflection(self, connection, metadata, partition_by):
        tn = "test_starrocks_reflection_partition_by_columns"
        connection.execute(text(f"DROP TABLE IF EXISTS {tn}"))
        create_table_sql = dedent(f"""
            CREATE TABLE {tn} (
                id INT,
                data VARCHAR(10),
                something_else DATETIME
            )
            PARTITION BY ({', '.join(partition_by)})
        """).strip()
        connection.execute(text(create_table_sql))

        reflected_meta = MetaData()
        reflected_meta.reflect(connection, only=[tn])
        reflected_table = reflected_meta.tables[tn]

        eq_(reflected_table.dialect_options["starrocks"]["partition_by"], ', '.join([f'`{c}`' for c in partition_by]))

    def test_expression_partitioning_reflection(self, connection, metadata):
        tn = "test_starrocks_reflection_expression_partitioning"
        connection.execute(text(f"DROP TABLE IF EXISTS {tn}"))
        create_table_sql = dedent(f"""
            CREATE TABLE {tn} (
                event_day DATETIME NOT NULL,
                data VARCHAR(10)
            )
            PARTITION BY date_trunc('month', event_day)
        """).strip()
        connection.execute(text(create_table_sql))

        reflected_meta = MetaData()
        reflected_meta.reflect(connection, only=[tn])
        reflected_table = reflected_meta.tables[tn]

        eq_(reflected_table.dialect_options["starrocks"]["partition_by"], "`event_day`")

    @testing.combinations(["id"], ["id", "data"], argnames="order_by")
    def test_order_by_reflection(self, connection, metadata, order_by):
        tn = "test_starrocks_reflection_order_by"
        connection.execute(text(f"DROP TABLE IF EXISTS {tn}"))
        create_table_sql = dedent(f"""
            CREATE TABLE {tn} (
                id INT,
                data VARCHAR(10),
                something_else DATETIME
            )
            ORDER BY ({', '.join(order_by)})
        """).strip()
        connection.execute(text(create_table_sql))

        reflected_meta = MetaData()
        reflected_meta.reflect(connection, only=[tn])
        reflected_table = reflected_meta.tables[tn]

        eq_(reflected_table.dialect_options["starrocks"]["order_by"], ', '.join([f'`{c}`' for c in order_by]))

    def test_nullable_column_reflection(self, connection, metadata):
        tn = "test_starrocks_reflection_nullable"
        connection.execute(text(f"DROP TABLE IF EXISTS {tn}"))
        create_table_sql = dedent(f"""
            CREATE TABLE {tn} (
                id INT NOT NULL,
                data VARCHAR(10) NULL,
                something_else DATETIME
            )
        """).strip()
        connection.execute(text(create_table_sql))

        reflected_meta = MetaData()
        reflected_meta.reflect(connection, only=[tn])
        reflected_table = reflected_meta.tables[tn]

        eq_(reflected_table.columns["id"].nullable, False)
        eq_(reflected_table.columns["data"].nullable, True)

    def test_generated_columns_reflection(self, connection, metadata):
        tn = "test_starrocks_reflection_generated"
        connection.execute(text(f"DROP TABLE IF EXISTS {tn}"))
        create_table_sql = dedent(f"""
            CREATE TABLE {tn} (
                id INT,
                data JSON,
                generated_column BIGINT AS (id + 1),
                generated_json TEXT AS get_json_string(`data`, '$.some_key') COMMENT 'generated from json'
            )
        """).strip()
        connection.execute(text(create_table_sql))

        reflected_meta = MetaData()
        reflected_meta.reflect(connection, only=[tn])
        reflected_table = reflected_meta.tables[tn]

        eq_(reflected_table.columns["generated_column"].computed.sqltext.text, "id + 1")
        eq_(reflected_table.columns["generated_json"].computed.sqltext.text, "get_json_string(`data`, '$.some_key')")
