from src.libs.custom_logger import get_custom_logger
import math
from time import time
import pytest
import allure
from src.libs.common import to_base64, delay
from src.data_classes import message_rpc_response_schema
from src.env_vars import NODE_1, NODE_2, NODEKEY
from src.node.waku_node import WakuNode
from tenacity import retry, stop_after_delay, wait_fixed

logger = get_custom_logger(__name__)


class StepsRelay:
    @pytest.fixture(scope="function", autouse=True)
    def setup_nodes(self, request):
        self.node1 = WakuNode(NODE_1, "node1_" + request.cls.test_id)
        self.node1.start(relay="true", discv5_discovery="true", peer_exchange="true", nodekey=NODEKEY)
        enr_uri = self.node1.info()["enrUri"]
        self.node2 = WakuNode(NODE_2, "node2_" + request.cls.test_id)
        self.node2.start(relay="true", discv5_discovery="true", discv5_bootstrap_node=enr_uri, peer_exchange="true")
        self.test_pubsub_topic = "/waku/2/rs/18/1"
        self.test_content_topic = "/test/1/waku-relay/proto"
        self.test_payload = "Relay works!!"
        self.node1.set_subscriptions([self.test_pubsub_topic])
        self.node2.set_subscriptions([self.test_pubsub_topic])

    @pytest.fixture(scope="function", autouse=True)
    def network_warm_up(self, setup_nodes):
        try:
            self.wait_for_published_message_to_reach_peer(120)
            logger.info("WARM UP successful !!")
        except Exception as ex:
            raise TimeoutError(f"WARM UP FAILED WITH: {ex}")

    @allure.step
    def check_published_message_reaches_peer(self, message, pubsub_topic=None, message_propagation_delay=0.1):
        self.node1.send_message(message, pubsub_topic or self.test_pubsub_topic)
        delay(message_propagation_delay)
        get_messages_response = self.node2.get_messages(pubsub_topic or self.test_pubsub_topic)
        assert get_messages_response, "Peer node couldn't find any messages"
        received_message = message_rpc_response_schema.load(get_messages_response[0])
        self.assert_received_message(message, received_message)

    def assert_received_message(self, sent_message, received_message):
        def assert_fail_message(field_name):
            return f"Incorrect field: {field_name}. Published: {sent_message[field_name]} Received: {getattr(received_message, field_name)}"

        assert received_message.payload == sent_message["payload"], assert_fail_message("payload")
        assert received_message.contentTopic == sent_message["contentTopic"], assert_fail_message("contentTopic")
        if sent_message.get("timestamp") is not None:
            if isinstance(sent_message["timestamp"], float):
                assert math.isclose(float(received_message.timestamp), sent_message["timestamp"], rel_tol=1e-9), assert_fail_message("timestamp")
            else:
                assert str(received_message.timestamp) == str(sent_message["timestamp"]), assert_fail_message("timestamp")
        if "version" in sent_message:
            assert str(received_message.version) == str(sent_message["version"]), assert_fail_message("version")
        if "meta" in sent_message:
            assert str(received_message.meta) == str(sent_message["meta"]), assert_fail_message("meta")
        if "ephemeral" in sent_message:
            assert str(received_message.ephemeral) == str(sent_message["ephemeral"]), assert_fail_message("ephemeral")
        if "rateLimitProof" in sent_message:
            assert str(received_message.rateLimitProof) == str(sent_message["rateLimitProof"]), assert_fail_message("rateLimitProof")

    def wait_for_published_message_to_reach_peer(self, timeout_duration, time_between_retries=1):
        @retry(stop=stop_after_delay(timeout_duration), wait=wait_fixed(time_between_retries), reraise=True)
        def check_peer_connection():
            message = {"payload": to_base64(self.test_payload), "contentTopic": self.test_content_topic, "timestamp": int(time() * 1e9)}
            self.check_published_message_reaches_peer(message)

        check_peer_connection()

    def ensure_subscriptions_on_nodes(self, node_list, pubsub_topic_list):
        for node in node_list:
            node.set_subscriptions(pubsub_topic_list)

    def create_message(self, **kwargs):
        message = {"payload": to_base64(self.test_payload), "contentTopic": self.test_content_topic, "timestamp": int(time() * 1e9)}
        message.update(kwargs)
        return message