import collections
import os
import re
import shutil
import time
import xml.etree.ElementTree as ET

import helpers.client
import helpers.cluster
from helpers.test_tools import TSV
import pytest


cluster = helpers.cluster.ClickHouseCluster(__file__)

node1 = cluster.add_instance(
    "node1",
    main_configs=["configs/logs_config.xml"],
    with_zookeeper=True,
    stay_alive=True,
    tmpfs=[
        "/jbod1:size=40M",
        "/jbod2:size=40M",
        "/jbod3:size=40M",
        "/jbod4:size=40M",
        "/external:size=200M",
    ],
    macros={"shard": 0, "replica": 1},
)

node2 = cluster.add_instance(
    "node2",
    main_configs=["configs/logs_config.xml"],
    with_zookeeper=True,
    stay_alive=True,
    tmpfs=[
        "/jbod1:size=40M",
        "/jbod2:size=40M",
        "/jbod3:size=40M",
        "/jbod4:size=40M",
        "/external:size=200M",
    ],
    macros={"shard": 0, "replica": 2},
)


def get_log(node):
    return node.exec_in_container(
        ["bash", "-c", "cat /var/log/clickhouse-server/clickhouse-server.log"]
    )


@pytest.fixture(scope="module")
def started_cluster():
    try:
        cluster.start()
        yield cluster

    finally:
        cluster.shutdown()


def start_over():
    shutil.copy(
        os.path.join(
            os.path.dirname(__file__), "configs/config.d/storage_configuration.xml"
        ),
        os.path.join(node1.config_d_dir, "storage_configuration.xml"),
    )

    for node in (node1, node2):
        separate_configuration_path = os.path.join(
            node.config_d_dir, "separate_configuration.xml"
        )
        try:
            os.remove(separate_configuration_path)
        except:
            """"""


def add_disk(node, name, path, separate_file=False):
    separate_configuration_path = os.path.join(
        node.config_d_dir, "separate_configuration.xml"
    )

    try:
        if separate_file:
            tree = ET.parse(separate_configuration_path)
        else:
            tree = ET.parse(
                os.path.join(node.config_d_dir, "storage_configuration.xml")
            )
    except:
        tree = ET.ElementTree(
            ET.fromstring(
                "<clickhouse><storage_configuration><disks/><policies/></storage_configuration></clickhouse>"
            )
        )
    root = tree.getroot()
    new_disk = ET.Element(name)
    new_path = ET.Element("path")
    new_path.text = path
    new_disk.append(new_path)
    root.find("storage_configuration").find("disks").append(new_disk)
    if separate_file:
        tree.write(separate_configuration_path)
    else:
        tree.write(os.path.join(node.config_d_dir, "storage_configuration.xml"))


def update_disk(node, name, path, keep_free_space_bytes, separate_file=False):
    separate_configuration_path = os.path.join(
        node.config_d_dir, "separate_configuration.xml"
    )

    try:
        if separate_file:
            tree = ET.parse(separate_configuration_path)
        else:
            tree = ET.parse(
                os.path.join(node.config_d_dir, "storage_configuration.xml")
            )
    except:
        tree = ET.ElementTree(
            ET.fromstring(
                "<clickhouse><storage_configuration><disks/><policies/></storage_configuration></clickhouse>"
            )
        )

    root = tree.getroot()
    disk = root.find("storage_configuration").find("disks").find(name)
    assert disk is not None

    new_path = disk.find("path")
    assert new_path is not None
    new_path.text = path

    new_keep_free_space_bytes = disk.find("keep_free_space_bytes")
    assert new_keep_free_space_bytes is not None
    new_keep_free_space_bytes.text = keep_free_space_bytes

    if separate_file:
        tree.write(separate_configuration_path)
    else:
        tree.write(os.path.join(node.config_d_dir, "storage_configuration.xml"))


def add_policy(node, name, volumes):
    tree = ET.parse(os.path.join(node.config_d_dir, "storage_configuration.xml"))
    root = tree.getroot()
    new_policy = ET.Element(name)
    new_volumes = ET.Element("volumes")
    for volume, disks in list(volumes.items()):
        new_volume = ET.Element(volume)
        for disk in disks:
            new_disk = ET.Element("disk")
            new_disk.text = disk
            new_volume.append(new_disk)
        new_volumes.append(new_volume)
    new_policy.append(new_volumes)
    root.find("storage_configuration").find("policies").append(new_policy)
    tree.write(os.path.join(node.config_d_dir, "storage_configuration.xml"))


def test_add_disk(started_cluster):
    try:
        name = "test_add_disk"
        engine = "MergeTree()"

        start_over()
        node1.restart_clickhouse(kill=True)
        time.sleep(2)

        node1.query(
            """
            CREATE TABLE {name} (
                d UInt64
            ) ENGINE = {engine}
            ORDER BY d
            SETTINGS storage_policy='jbods_with_external'
        """.format(
                name=name, engine=engine
            )
        )

        assert "jbod3" not in set(
            node1.query("SELECT name FROM system.disks").splitlines()
        )

        add_disk(node1, "jbod3", "/jbod3/")
        node1.query("SYSTEM RELOAD CONFIG")

        assert "jbod3" in set(node1.query("SELECT name FROM system.disks").splitlines())
    finally:
        try:
            node1.query("DROP TABLE IF EXISTS {}".format(name))
        except:
            """"""


def test_update_disk(started_cluster):
    try:
        name = "test_update_disk"
        engine = "MergeTree()"

        start_over()
        node1.restart_clickhouse(kill=True)
        time.sleep(2)

        node1.query(
            """
            CREATE TABLE {name} (
                d UInt64
            ) ENGINE = {engine}
            ORDER BY d
            SETTINGS storage_policy='jbods_with_external'
        """.format(
                name=name, engine=engine
            )
        )

        assert node1.query(
            "SELECT path, keep_free_space FROM system.disks where name = 'jbod2'"
        ) == TSV([["/jbod2/", "10485760"]])

        update_disk(node1, "jbod2", "/jbod2/", "20971520")
        node1.query("SYSTEM RELOAD CONFIG")

        assert node1.query(
            "SELECT path, keep_free_space FROM system.disks where name = 'jbod2'"
        ) == TSV([["/jbod2/", "20971520"]])
    finally:
        try:
            node1.query("DROP TABLE IF EXISTS {}".format(name))
        except:
            """"""


def test_add_disk_to_separate_config(started_cluster):
    try:
        name = "test_add_disk"
        engine = "MergeTree()"

        start_over()
        node1.restart_clickhouse(kill=True)
        time.sleep(2)

        node1.query(
            """
            CREATE TABLE {name} (
                d UInt64
            ) ENGINE = {engine}
            ORDER BY d
            SETTINGS storage_policy='jbods_with_external'
        """.format(
                name=name, engine=engine
            )
        )

        assert "jbod3" not in set(
            node1.query("SELECT name FROM system.disks").splitlines()
        )

        add_disk(node1, "jbod3", "/jbod3/", separate_file=True)
        node1.query("SYSTEM RELOAD CONFIG")

        assert "jbod3" in set(node1.query("SELECT name FROM system.disks").splitlines())
        start_over()

    finally:
        try:
            node1.query("DROP TABLE IF EXISTS {}".format(name))
        except:
            """"""


def test_add_policy(started_cluster):
    try:
        name = "test_add_policy"
        engine = "MergeTree()"

        start_over()
        add_disk(node1, "jbod3", "/jbod3/")
        add_disk(node1, "jbod4", "/jbod4/")
        node1.restart_clickhouse(kill=True)
        time.sleep(2)

        node1.query(
            """
            CREATE TABLE {name} (
                d UInt64
            ) ENGINE = {engine}
            ORDER BY d
            SETTINGS storage_policy='jbods_with_external'
        """.format(
                name=name, engine=engine
            )
        )

        add_policy(node1, "cool_policy", {"volume1": ["jbod3", "jbod4"]})
        node1.query("SYSTEM RELOAD CONFIG")

        disks = set(node1.query("SELECT name FROM system.disks").splitlines())
        assert "cool_policy" in set(
            node1.query("SELECT policy_name FROM system.storage_policies").splitlines()
        )
        assert {"volume1"} == set(
            node1.query(
                "SELECT volume_name FROM system.storage_policies WHERE policy_name = 'cool_policy'"
            ).splitlines()
        )
        assert {"['jbod3','jbod4']"} == set(
            node1.query(
                "SELECT disks FROM system.storage_policies WHERE policy_name = 'cool_policy'"
            ).splitlines()
        )

    finally:
        try:
            node1.query("DROP TABLE IF EXISTS {}".format(name))
        except:
            """"""


def test_new_policy_works(started_cluster):
    try:
        name = "test_new_policy_works"
        engine = "MergeTree()"

        start_over()
        add_disk(node1, "jbod3", "/jbod3/")
        node1.restart_clickhouse(kill=True)
        time.sleep(2)

        node1.query(
            """
            CREATE TABLE {name} (
                d UInt64
            ) ENGINE = {engine}
            ORDER BY d
            SETTINGS storage_policy='jbods_with_external'
        """.format(
                name=name, engine=engine
            )
        )

        add_policy(node1, "cool_policy", {"volume1": ["jbod3"]})
        node1.query("SYSTEM RELOAD CONFIG")

        # Incompatible storage policy.
        with pytest.raises(helpers.client.QueryRuntimeException):
            node1.query(
                """
                ALTER TABLE {name} MODIFY SETTING storage_policy='cool_policy'
            """.format(
                    name=name
                )
            )

        start_over()
        add_disk(node1, "jbod3", "/jbod3/")
        add_disk(node1, "jbod4", "/jbod4/")
        add_policy(
            node1,
            "cool_policy",
            collections.OrderedDict(
                [
                    ("volume1", ["jbod3"]),
                    ("main", ["jbod1", "jbod2"]),
                    ("external", ["external"]),
                ]
            ),
        )
        node1.query("SYSTEM RELOAD CONFIG")

        node1.query(
            """
            ALTER TABLE {name} MODIFY SETTING storage_policy='cool_policy'
        """.format(
                name=name
            )
        )

        node1.query(
            """
            INSERT INTO TABLE {name} VALUES (1)
        """.format(
                name=name
            )
        )
        assert {"jbod3"} == set(
            node1.query(
                "SELECT disk_name FROM system.parts WHERE active = 1 AND table = '{name}'".format(
                    name=name
                )
            ).splitlines()
        )

    finally:
        try:
            node1.query("DROP TABLE IF EXISTS {}".format(name))
        except:
            """"""


def test_add_volume_to_policy(started_cluster):
    try:
        name = "test_add_volume_to_policy"
        engine = "MergeTree()"

        start_over()
        add_disk(node1, "jbod3", "/jbod3/")
        add_disk(node1, "jbod4", "/jbod4/")
        add_policy(node1, "cool_policy", {"volume1": ["jbod3"]})
        node1.restart_clickhouse(kill=True)
        time.sleep(2)

        node1.query(
            """
            CREATE TABLE {name} (
                d UInt64
            ) ENGINE = {engine}
            ORDER BY d
            SETTINGS storage_policy='jbods_with_external'
        """.format(
                name=name, engine=engine
            )
        )

        start_over()
        add_disk(node1, "jbod3", "/jbod3/")
        add_disk(node1, "jbod4", "/jbod4/")
        add_policy(
            node1,
            "cool_policy",
            collections.OrderedDict([("volume1", ["jbod3"]), ("volume2", ["jbod4"])]),
        )
        node1.query("SYSTEM RELOAD CONFIG")

        volumes = set(
            node1.query(
                "SELECT volume_name FROM system.storage_policies WHERE policy_name = 'cool_policy'"
            ).splitlines()
        )
        disks_sets = set(
            node1.query(
                "SELECT disks FROM system.storage_policies WHERE policy_name = 'cool_policy'"
            ).splitlines()
        )
        assert {"volume1", "volume2"} == volumes
        assert {"['jbod3']", "['jbod4']"} == disks_sets

    finally:
        try:
            node1.query("DROP TABLE IF EXISTS {}".format(name))
        except:
            """"""


def test_add_disk_to_policy(started_cluster):
    try:
        name = "test_add_disk_to_policy"
        engine = "MergeTree()"

        start_over()
        add_disk(node1, "jbod3", "/jbod3/")
        add_disk(node1, "jbod4", "/jbod4/")
        add_policy(node1, "cool_policy", {"volume1": ["jbod3"]})
        node1.restart_clickhouse(kill=True)
        time.sleep(2)

        node1.query(
            """
            CREATE TABLE {name} (
                d UInt64
            ) ENGINE = {engine}
            ORDER BY d
            SETTINGS storage_policy='jbods_with_external'
        """.format(
                name=name, engine=engine
            )
        )

        start_over()
        add_disk(node1, "jbod3", "/jbod3/")
        add_disk(node1, "jbod4", "/jbod4/")
        add_policy(node1, "cool_policy", {"volume1": ["jbod3", "jbod4"]})
        node1.query("SYSTEM RELOAD CONFIG")

        volumes = set(
            node1.query(
                "SELECT volume_name FROM system.storage_policies WHERE policy_name = 'cool_policy'"
            ).splitlines()
        )
        disks_sets = set(
            node1.query(
                "SELECT disks FROM system.storage_policies WHERE policy_name = 'cool_policy'"
            ).splitlines()
        )
        assert {"volume1"} == volumes
        assert {"['jbod3','jbod4']"} == disks_sets

    finally:
        try:
            node1.query("DROP TABLE IF EXISTS {}".format(name))
        except:
            """"""


def test_remove_disk(started_cluster):
    try:
        name = "test_remove_disk"
        engine = "MergeTree()"

        start_over()
        add_disk(node1, "remove_disk_jbod3", "/jbod3/")
        node1.restart_clickhouse(kill=True)
        time.sleep(2)

        node1.query(
            """
            CREATE TABLE {name} (
                d UInt64
            ) ENGINE = {engine}
            ORDER BY d
            SETTINGS storage_policy='jbods_with_external'
        """.format(
                name=name, engine=engine
            )
        )

        assert "remove_disk_jbod3" in set(
            node1.query("SELECT name FROM system.disks").splitlines()
        )

        start_over()
        node1.query("SYSTEM RELOAD CONFIG")

        assert "remove_disk_jbod3" in set(
            node1.query("SELECT name FROM system.disks").splitlines()
        )
        assert re.search("Warning.*remove_disk_jbod3", get_log(node1))
    finally:
        try:
            node1.query("DROP TABLE IF EXISTS {}".format(name))
        except:
            """"""


def test_remove_policy(started_cluster):
    try:
        name = "test_remove_policy"
        engine = "MergeTree()"

        start_over()
        add_disk(node1, "jbod3", "/jbod3/")
        add_disk(node1, "jbod4", "/jbod4/")
        add_policy(node1, "remove_policy_cool_policy", {"volume1": ["jbod3", "jbod4"]})
        node1.restart_clickhouse(kill=True)
        time.sleep(2)

        node1.query(
            """
            CREATE TABLE {name} (
                d UInt64
            ) ENGINE = {engine}
            ORDER BY d
            SETTINGS storage_policy='jbods_with_external'
        """.format(
                name=name, engine=engine
            )
        )

        assert "remove_policy_cool_policy" in set(
            node1.query("SELECT policy_name FROM system.storage_policies").splitlines()
        )

        start_over()
        add_disk(node1, "jbod3", "/jbod3/")
        add_disk(node1, "jbod4", "/jbod4/")
        node1.query("SYSTEM RELOAD CONFIG")

        assert "remove_policy_cool_policy" in set(
            node1.query("SELECT policy_name FROM system.storage_policies").splitlines()
        )
        assert re.search("Error.*remove_policy_cool_policy", get_log(node1))

    finally:
        try:
            node1.query("DROP TABLE IF EXISTS {}".format(name))
        except:
            """"""


def test_remove_volume_from_policy(started_cluster):
    try:
        name = "test_remove_volume_from_policy"
        engine = "MergeTree()"

        start_over()
        add_disk(node1, "jbod3", "/jbod3/")
        add_disk(node1, "jbod4", "/jbod4/")
        add_policy(
            node1,
            "test_remove_volume_from_policy_cool_policy",
            collections.OrderedDict([("volume1", ["jbod3"]), ("volume2", ["jbod4"])]),
        )
        node1.restart_clickhouse(kill=True)
        time.sleep(2)

        node1.query(
            """
            CREATE TABLE {name} (
                d UInt64
            ) ENGINE = {engine}
            ORDER BY d
            SETTINGS storage_policy='jbods_with_external'
        """.format(
                name=name, engine=engine
            )
        )

        volumes = set(
            node1.query(
                "SELECT volume_name FROM system.storage_policies WHERE policy_name = 'test_remove_volume_from_policy_cool_policy'"
            ).splitlines()
        )
        disks_sets = set(
            node1.query(
                "SELECT disks FROM system.storage_policies WHERE policy_name = 'test_remove_volume_from_policy_cool_policy'"
            ).splitlines()
        )
        assert {"volume1", "volume2"} == volumes
        assert {"['jbod3']", "['jbod4']"} == disks_sets

        start_over()
        add_disk(node1, "jbod3", "/jbod3/")
        add_disk(node1, "jbod4", "/jbod4/")
        add_policy(node1, "cool_policy", {"volume1": ["jbod3"]})
        node1.query("SYSTEM RELOAD CONFIG")

        volumes = set(
            node1.query(
                "SELECT volume_name FROM system.storage_policies WHERE policy_name = 'test_remove_volume_from_policy_cool_policy'"
            ).splitlines()
        )
        disks_sets = set(
            node1.query(
                "SELECT disks FROM system.storage_policies WHERE policy_name = 'test_remove_volume_from_policy_cool_policy'"
            ).splitlines()
        )
        assert {"volume1", "volume2"} == volumes
        assert {"['jbod3']", "['jbod4']"} == disks_sets
        assert re.search(
            "Error.*test_remove_volume_from_policy_cool_policy", get_log(node1)
        )

    finally:
        try:
            node1.query("DROP TABLE IF EXISTS {}".format(name))
        except:
            """"""


def test_remove_disk_from_policy(started_cluster):
    try:
        name = "test_remove_disk_from_policy"
        engine = "MergeTree()"

        start_over()
        add_disk(node1, "jbod3", "/jbod3/")
        add_disk(node1, "jbod4", "/jbod4/")
        add_policy(
            node1,
            "test_remove_disk_from_policy_cool_policy",
            {"volume1": ["jbod3", "jbod4"]},
        )
        node1.restart_clickhouse(kill=True)
        time.sleep(2)

        node1.query(
            """
            CREATE TABLE {name} (
                d UInt64
            ) ENGINE = {engine}
            ORDER BY d
            SETTINGS storage_policy='jbods_with_external'
        """.format(
                name=name, engine=engine
            )
        )

        volumes = set(
            node1.query(
                "SELECT volume_name FROM system.storage_policies WHERE policy_name = 'test_remove_disk_from_policy_cool_policy'"
            ).splitlines()
        )
        disks_sets = set(
            node1.query(
                "SELECT disks FROM system.storage_policies WHERE policy_name = 'test_remove_disk_from_policy_cool_policy'"
            ).splitlines()
        )
        assert {"volume1"} == volumes
        assert {"['jbod3','jbod4']"} == disks_sets

        start_over()
        add_disk(node1, "jbod3", "/jbod3/")
        add_disk(node1, "jbod4", "/jbod4/")
        add_policy(node1, "cool_policy", {"volume1": ["jbod3"]})
        node1.query("SYSTEM RELOAD CONFIG")

        volumes = set(
            node1.query(
                "SELECT volume_name FROM system.storage_policies WHERE policy_name = 'test_remove_disk_from_policy_cool_policy'"
            ).splitlines()
        )
        disks_sets = set(
            node1.query(
                "SELECT disks FROM system.storage_policies WHERE policy_name = 'test_remove_disk_from_policy_cool_policy'"
            ).splitlines()
        )
        assert {"volume1"} == volumes
        assert {"['jbod3','jbod4']"} == disks_sets
        assert re.search(
            "Error.*test_remove_disk_from_policy_cool_policy", get_log(node1)
        )

    finally:
        try:
            node1.query("DROP TABLE IF EXISTS {}".format(name))
        except:
            """"""