#!/usr/bin/env python3
import os
from random import randint, choices
import sys

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()

N = 10
create_query = (
    "CREATE TABLE t_cnf_fuzz("
    + ", ".join([f"c{i} UInt8" for i in range(N)])
    + ") ENGINE = Memory"
)

client.query("DROP TABLE IF EXISTS t_cnf_fuzz")
client.query(create_query)

# Insert all possible combinations of bool columns.
insert_query = "INSERT INTO t_cnf_fuzz VALUES "
for i in range(2**N):
    values = []
    cur = i
    for _ in range(N):
        values.append(cur % 2)
        cur //= 2

    insert_query += "(" + ", ".join(map(lambda x: str(x), values)) + ")"

client.query(insert_query)

# Let's try to covert DNF to CNF,
# because it's a worst case in a sense.

MAX_CLAUSES = 10
MAX_ATOMS = 5


def generate_dnf():
    clauses = []
    num_clauses = randint(1, MAX_CLAUSES)
    for _ in range(num_clauses):
        num_atoms = randint(1, MAX_ATOMS)
        atom_ids = choices(range(N), k=num_atoms)
        negates = choices([0, 1], k=num_atoms)
        atoms = [
            f"(NOT c{i})" if neg else f"c{i}" for (i, neg) in zip(atom_ids, negates)
        ]
        clauses.append("(" + " AND ".join(atoms) + ")")

    return " OR ".join(clauses)


select_query = (
    "SELECT count() FROM t_cnf_fuzz WHERE {} SETTINGS convert_query_to_cnf = {}"
)

fail_report = """
Failed query: '{}'.
Result without optimization: {}.
Result with optimization: {}.
"""

T = 500
for _ in range(T):
    condition = generate_dnf()

    query = select_query.format(condition, 0)
    res = client.query(query).strip()

    query_cnf = select_query.format(condition, 1)
    res_cnf = client.query(query_cnf).strip()

    if res != res_cnf:
        print(fail_report.format(query_cnf, res, res_cnf))
        exit(1)

client.query("DROP TABLE t_cnf_fuzz")
print("OK")