from valkey_search_test_case import ValkeySearchClusterTestCaseDebugMode
from valkey.cluster import ValkeyCluster
from valkey.client import Valkey
from valkeytestframework.conftest import resource_port_tracker
from valkey.exceptions import ResponseError
from indexes import *
from test_cancel import Any
from valkeytestframework.util import waiters
from valkey import ResponseError
from test_fanout_base import MAX_RETRIES
from ft_info_parser import FTInfoParser

class TestFTInfoPartitionConsistencyControls(ClusterTestUtils, ValkeySearchClusterTestCaseDebugMode):
    def run_info_command(self, client, index_name, enable_partial_results=False, require_consistency=False, expect_error=False):
        try:
            return client.execute_command(
                "FT.INFO",
                index_name,
                "PRIMARY",
                "SOMESHARDS" if enable_partial_results else "ALLSHARDS",
                "CONSISTENT" if require_consistency else "INCONSISTENT"
            )
        except ResponseError as e:
            if expect_error:
                assert str(e) == "Unable to contact all cluster members"
                return []
            else:
                raise

    def test_ft_info_consistency_controls(self):
        self.execute_primaries(["flushall sync"])
        self.config_set("search.info-developer-visible", "yes")
        client: Valkey = self.new_cluster_client()

        index_name = "hnsw"
        index = Index(index_name, [Vector("v", 3, type="HNSW"), Numeric("n")])
        index.create(client)
        index.load_data(client, 1000)
        assert self.sum_docs(index) == 1000

        # normal result without consistency check
        normal_result = self.run_info_command(client, index_name, require_consistency=False)

        # normal result with consistency check
        cur_result = self.run_info_command(client, index_name, require_consistency=True)
        assert FTInfoParser(cur_result).num_docs == FTInfoParser(normal_result).num_docs
        
        # force invalid invalid index fingerprint and version
        self.control_set("ForceInfoInvalidSlotFingerprint", "yes")
        self.config_set("search.ft-info-timeout-ms", "500")

        # enable consistency check, get error result
        cur_result = self.run_info_command(client, index_name, require_consistency=True, expect_error=True)
        assert cur_result == []

        # disable consistency check, get normal result
        cur_result = self.run_info_command(client, index_name, require_consistency=False, expect_error=False)

        assert FTInfoParser(cur_result).num_docs == FTInfoParser(normal_result).num_docs

        self.control_set("ForceInfoInvalidSlotFingerprint", "no")

    def test_ft_info_partition_controls(self):
        self.execute_primaries(["flushall sync"])
        self.config_set("search.info-developer-visible", "yes")
        client: Valkey = self.new_cluster_client()

        index_name = "hnsw"
        index = Index(index_name, [Vector("v", 3, type="HNSW"), Numeric("n")])
        index.create(client)
        index.load_data(client, 1000)
        assert self.sum_docs(index) == 1000

        # normal result without partition controls
        normal_result = self.run_info_command(client, index_name)

        # normal result with partition controls
        cur_result = self.run_info_command(client, index_name, enable_partial_results=True)
        assert FTInfoParser(cur_result).num_docs == FTInfoParser(normal_result).num_docs

        self.config_set("search.ft-info-timeout-ms", "500")

        node0: Valkey = self.new_client_for_primary(0)
        node1: Valkey = self.new_client_for_primary(1)
        
        # force timeout on a remote node
        assert node1.execute_command(
            "FT._DEBUG CONTROLLED_VARIABLE SET ForceRemoteFailCount ", 
            MAX_RETRIES
        ) == b"OK"

        # disable partial result, get error result
        cur_result = self.run_info_command(node0, index_name, enable_partial_results=False, expect_error=True)
        assert cur_result == []

        # enable partial results, get partial result
        cur_result = self.run_info_command(node0, index_name, enable_partial_results=True)
        assert FTInfoParser(cur_result).num_docs != FTInfoParser(normal_result).num_docs

        # reset timeout on remote node
        assert node1.execute_command(
            "FT._DEBUG CONTROLLED_VARIABLE SET ForceRemoteFailCount ", 
            0
        ) == b"OK"
