import time
import pytest

import redis
from redis.client import Script
from redis.exceptions import ResponseError

from rb.cluster import Cluster
from rb.router import UnroutableCommand
from rb.promise import Promise
from rb.utils import text_type


def test_basic_interface():
    cluster = Cluster(
        {0: {"db": 0}, 1: {"db": 2}, 2: {"db": 4, "host": "127.0.0.1"},},
        host_defaults={"password": "pass",},
        pool_options={"encoding": "utf-8", "decode_responses": True},
    )

    assert len(cluster.hosts) == 3

    assert cluster.hosts[0].host_id == 0
    assert cluster.hosts[0].db == 0
    assert cluster.hosts[0].host == "localhost"
    assert cluster.hosts[0].port == 6379
    assert cluster.hosts[0].password == "pass"

    assert cluster.hosts[1].host_id == 1
    assert cluster.hosts[1].db == 2
    assert cluster.hosts[1].host == "localhost"
    assert cluster.hosts[1].port == 6379
    assert cluster.hosts[1].password == "pass"

    assert cluster.hosts[2].host_id == 2
    assert cluster.hosts[2].db == 4
    assert cluster.hosts[2].host == "127.0.0.1"
    assert cluster.hosts[2].port == 6379
    assert cluster.hosts[2].password == "pass"


def test_router_access():
    cluster = Cluster(
        {0: {"db": 0},}, pool_options={"encoding": "utf-8", "decode_responses": True}
    )

    router = cluster.get_router()
    assert router.cluster is cluster
    assert cluster.get_router() is router

    cluster.add_host(1, {"db": 1})
    new_router = cluster.get_router()
    assert new_router is not router


def test_basic_cluster(cluster):
    iterations = 10000

    with cluster.map() as client:
        for x in range(iterations):
            client.set("key-%06d" % x, x)
    responses = []
    with cluster.map() as client:
        for x in range(iterations):
            responses.append(client.get("key-%06d" % x))
    ref_sum = sum(int(x.value) for x in responses)
    assert ref_sum == sum(range(iterations))


def test_basic_cluster_disabled_batch(cluster):
    iterations = 10000

    with cluster.map(auto_batch=False) as client:
        for x in range(iterations):
            client.set("key-%06d" % x, x)
    responses = []
    with cluster.map(auto_batch=False) as client:
        for x in range(iterations):
            responses.append(client.get("key-%06d" % x))
    ref_sum = sum(int(x.value) for x in responses)
    assert ref_sum == sum(range(iterations))


def make_zset_data(x):
    return [(str(i), float(i)) for i in range(x, x + 10)]


def test_simple_api(cluster):
    client = cluster.get_routing_client()
    with client.map() as map_client:
        for x in range(10):
            map_client.set("key:%d" % x, x)
            if redis.VERSION >= (3, 0, 0):
                map_client.zadd("zset:%d" % x, dict(make_zset_data(x)))
            else:
                map_client.zadd("zset:%d" % x, **dict(make_zset_data(x)))

    for x in range(10):
        assert client.get("key:%d" % x) == str(x)
        assert client.zrange("zset:%d" % x, 0, -1, withscores=True) == make_zset_data(x)

    results = []  # (promise, expected result)
    with client.map() as map_client:
        for x in range(10):
            results.append(
                (
                    map_client.zrange("zset:%d" % x, 0, -1, withscores=True),
                    make_zset_data(x),
                )
            )

    for promise, expectation in results:
        assert promise.value == expectation

    with client.map() as map_client:
        for x in range(10):
            map_client.delete("key:%d" % x)

    for x in range(10):
        assert client.get("key:%d" % x) is None


def test_routing_client_releases_connection_on_error(cluster):
    client = cluster.get_routing_client()
    with pytest.raises(ResponseError):
        client.sadd("key")

    host = cluster.get_router().get_host_for_command("sadd", ["key"])
    pool = cluster.get_pool_for_host(host)
    assert len(pool._available_connections) == pool._created_connections


def test_mapping_client_releases_connection_on_error(cluster):
    client = cluster.get_routing_client().get_mapping_client()
    client.sadd("key")
    with pytest.raises(ResponseError):
        client.join()

    host = cluster.get_router().get_host_for_command("sadd", ["key"])
    pool = cluster.get_pool_for_host(host)
    assert len(pool._available_connections) == pool._created_connections


def test_managed_mapping_client_releases_connection_on_error(cluster):
    with pytest.raises(ResponseError):
        with cluster.get_routing_client().map() as client:
            client.sadd("key")

    host = cluster.get_router().get_host_for_command("sadd", ["key"])
    pool = cluster.get_pool_for_host(host)
    assert len(pool._available_connections) == pool._created_connections


def test_multi_keys_rejected(cluster):
    client = cluster.get_routing_client()

    # Okay
    with client.map() as map_client:
        map_client.delete("key")