#!/usr/bin/env python3
import os
import sys
from itertools import chain, combinations, permutations

CURDIR = os.path.dirname(os.path.realpath(__file__))
sys.path.insert(0, os.path.join(CURDIR, "helpers"))

from pure_http_client import ClickHouseClient

client = ClickHouseClient()


def powerset(iterable):
    s = list(iterable)
    return chain.from_iterable(combinations(s, r) for r in range(len(s) + 1))


queries = [
    {"optimize": True, "where": [], "order_by": ["a"]},
    {"optimize": True, "where": [], "order_by": ["a", "b"]},
    {"optimize": True, "where": [], "order_by": ["a", "b", "c"]},
    {"optimize": True, "where": [], "order_by": ["a", "b", "c", "d"]},
    {"optimize": True, "where": ["a"], "order_by": ["a"]},
    {"optimize": True, "where": ["a"], "order_by": ["a", "b"]},
    {"optimize": True, "where": ["a"], "order_by": ["b"]},
    {"optimize": True, "where": ["a"], "order_by": ["b", "c"]},
    {"optimize": True, "where": ["b"], "order_by": ["a"]},
    {"optimize": True, "where": ["b"], "order_by": ["a", "c"]},
    {"optimize": False, "where": ["b"], "order_by": ["b", "c"]},
    {"optimize": True, "where": ["c"], "order_by": ["a"]},
    {"optimize": True, "where": ["c"], "order_by": ["a", "b"]},
    {"optimize": True, "where": ["a", "b"], "order_by": ["a"]},
    {"optimize": True, "where": ["a", "b"], "order_by": ["a", "b"]},
    {"optimize": True, "where": ["a", "b"], "order_by": ["a", "c"]},
    {"optimize": True, "where": ["a", "b"], "order_by": ["a", "b", "c"]},
    {"optimize": True, "where": ["a", "b"], "order_by": ["a", "b", "c", "d"]},
    {"optimize": True, "where": ["a", "b"], "order_by": ["b", "c"]},
    {"optimize": True, "where": ["a", "b"], "order_by": ["c", "d"]},
    {"optimize": True, "where": ["a", "c"], "order_by": ["a"]},
    {"optimize": True, "where": ["a", "c"], "order_by": ["a", "b"]},
    {"optimize": True, "where": ["a", "c"], "order_by": ["b", "d"]},
    {"optimize": True, "where": ["a", "c"], "order_by": ["a", "b", "c"]},
    {"optimize": True, "where": ["a", "c"], "order_by": ["b", "c", "d"]},
    {"optimize": True, "where": ["a", "c"], "order_by": ["a", "b", "c", "d"]},
    {"optimize": False, "where": [], "order_by": ["b"]},
    {"optimize": False, "where": [], "order_by": ["b", "a"]},
    {"optimize": False, "where": [], "order_by": ["b", "c"]},
    {"optimize": False, "where": ["a"], "order_by": ["c"]},
    {"optimize": False, "where": ["a"], "order_by": ["c", "b"]},
    {"optimize": False, "where": ["a"], "order_by": ["c", "d"]},
    {"optimize": False, "where": ["c"], "order_by": ["c", "d"]},
    {"optimize": False, "where": ["c"], "order_by": ["b", "c"]},
]


client.query("DROP TABLE IF EXISTS t_fixed_prefix")
client.query(
    """
CREATE TABLE t_fixed_prefix (a UInt32, b UInt32, c UInt32, d UInt32, e UInt32)
ENGINE = MergeTree ORDER BY (a, b, c, d)"""
)

client.query("SYSTEM STOP MERGES t_fixed_prefix")

# create several parts
for _ in range(4):
    client.query(
        "INSERT INTO t_fixed_prefix SELECT number % 2, number % 10, number % 100, number % 1000, number FROM numbers(25000)"
    )


def check_query(
    fixed_columns, order_by_columns, should_be_optimized, should_use_finish_sorting
):
    where_clause = " AND ".join([f"{c} = 1" for c in fixed_columns])
    order_by_clause = ", ".join(order_by_columns)

    query = "SELECT {} FROM t_fixed_prefix".format(order_by_clause)
    if len(where_clause) != 0:
        query += " WHERE " + where_clause

    if len(order_by_clause) != 0:
        query += " ORDER BY " + order_by_clause

    query += " SETTINGS optimize_read_in_order = {}"

    res_optimized = client.query(query.format(1))
    res_not_optimized = client.query(query.format(0))

    if res_optimized != res_not_optimized:
        print("Failed query {}. Result of queries mismatched".format(query))
        exit(1)

    res_explain = client.query("EXPLAIN PIPELINE {}".format(query.format(1)))

    is_optimized = "MergeSortingTransform" not in res_explain
    uses_finish_sorting = "FinishSortingTransform" in res_explain

    if (
        is_optimized != should_be_optimized
        or uses_finish_sorting != should_use_finish_sorting
    ):
        print(
            """
Wrong query pipeline is built for query {}:
{}
Should be optimized: {}.
Is optimized: {}.
Should use FinishSortingTransform: {}.
Uses FinishSortingTransform: {}
""".format(
                query.format(1),
                res_explain,
                should_be_optimized,
                is_optimized,
                should_use_finish_sorting,
                uses_finish_sorting,
            )
        )
        exit(1)


for query in queries:
    check_query(query["where"], query["order_by"], query["optimize"], False)
    check_query(
        query["where"], query["order_by"] + ["e"], query["optimize"], query["optimize"]
    )

    where_columns = [f"bitNot({col})" for col in query["where"]]
    check_query(where_columns, query["order_by"], query["optimize"], False)
    check_query(
        where_columns, query["order_by"] + ["e"], query["optimize"], query["optimize"]
    )

print("OK")
