from typing import Dict, List, Optional, cast
from harness.src.interface import Interface
from rraft import (
    Config,
    ConfState,
    ConfigRef,
    Logger,
    LoggerRef,
    MemStorage,
    Message,
    MessageRef,
    MessageType,
    NO_LIMIT,
    InMemoryRaft,
)
import random


class Connection:
    """
    A connection from one node to another.

    Used in by `Network` for determining drop rates on messages.
    """
    def __hash__(self) -> int:
        return hash((self.from_, self.to))

    def __eq__(self, other) -> bool:
        return self.from_ == other.from_ and self.to == other.to

    def __init__(self, from_: int, to: int) -> None:
        self.from_ = from_
        self.to = to


class Network:
    """
    A simulated network for testing.

    You can use this to create a test network of Raft nodes.

    *Please note:* no actual network calls are made.
    """
    def __init__(
        self,
        peers: Dict[int, Interface],
        storage: Dict[int, MemStorage],
        dropm: Dict[Connection, float],
        ignorem: Dict[MessageType, bool],
    ) -> None:
        self.peers = peers
        self.storage = storage
        self.dropm = dropm
        self.ignorem = ignorem

    @staticmethod
    def default_config() -> Config:
        """
        Get a base config. Calling `Network::new` will initialize peers with this config.
        """
        cfg = Config.default()
        cfg.set_election_tick(10)
        cfg.set_heartbeat_tick(1)
        cfg.set_max_size_per_msg(NO_LIMIT)
        cfg.set_max_inflight_msgs(256)
        return cfg

    @staticmethod
    def new(peers: List[Optional[Interface]], l: Logger | LoggerRef) -> "Network":
        """
        Initializes a network from `peers`.

        Nodes will receive their ID based on their index in the vector, starting with 1.

        A `None` node will be replaced with a new Raft node, and its configuration will
        be `peers`.
        """
        return Network.new_with_config(peers, Network.default_config(), l)

    @staticmethod
    def new_with_config(
        peers: List[Optional[Interface]],
        config: Config | ConfigRef,
        l: Logger | LoggerRef,
    ) -> "Network":
        """
        Initialize a network from `peers` with explicitly specified `config`.
        """
        nstorage = {}
        npeers = {}
        peer_addrs = list(range(1, len(peers) + 1))

        for p, id in zip(peers, peer_addrs):
            if p is None:
                store = MemStorage.new_with_conf_state(
                    ConfState(peer_addrs, [])
                )
                nstorage[id] = store.clone()
                cfg = config.clone()
                cfg.set_id(id)
                npeers[id] = Interface(InMemoryRaft(cfg, store, l))
            else:
                if p.raft:
                    if raft := p.raft:
                        assert (
                            raft.get_id() == id
                        ), f"peer {p.raft.get_id()} in peers has a wrong position"

                        store = raft.get_raft_log().get_store().clone()
                        nstorage[id] = store

                npeers[id] = p

        return Network(npeers, nstorage, {}, {})

    def ignore(self, t: MessageType) -> None:
        """
        Ignore a given `MessageType`.
        """
        self.ignorem[t] = True

    def filter_(
        self, msgs: List[Message] | List[MessageRef]
    ) -> List[Message]:
        """
        Filter out messages that should be dropped according to rules set by `ignore` or `drop`.
        """
        def should_be_filtered(m: Message | MessageRef):
            if self.ignorem.get(m.get_msg_type()):
                return False

            # hups never go over the network, so don't drop them but panic
            assert m.get_msg_type() != MessageType.MsgHup, "unexpected msgHup"

            perc = self.dropm.get(Connection(m.get_from(), m.get_to()), 0.0)

            return random.random() >= perc

        return list(filter(should_be_filtered, msgs))

    def read_messages(self) -> List[Message]:
        """
        # Read out all messages generated by peers in the `Network`.
        #
        # Note: messages are not filtered by any configured filters.
        """
        msgs = []
        for _id, p in self.peers.items():
            msgs.extend(p.read_messages())
        return msgs

    def send(self, msgs: List[Message] | List[MessageRef]) -> None:
        """
        # Instruct the cluster to `step` through the given messages.
        #
        # NOTE: the given `msgs` won't be filtered by its filters.
        """
        while msgs:
            new_msgs: List[Message] = []

            for m in msgs:
                p = cast(Interface, self.peers.get(m.get_to()))
                try:
                    p.step(m)
                except Exception:
                    continue
                finally:
                    # The unstable data should be persisted before sending msg.
                    p.persist()
                    resp = p.read_messages()
                    new_msgs.extend(self.filter_(resp))

            msgs = []
            msgs.extend(new_msgs)

    def filter_and_send(self, msgs: List[Message] | List[MessageRef]) -> None:
        """
        Filter `msgs` and then instruct the cluster to `step` through the given messages.
        """
        self.send(self.filter_(msgs))

    def dispatch(self, messages: List[Message] | List[MessageRef]) -> None:
        """
        # Dispatches the given messages to the appropriate peers.
        #
        # Unlike `send` this does not gather and send any responses. It also does not ignore errors.]
        """
        for message in self.filter_(messages):
            to = message.get_to()
            peer = self.peers[to]
            peer.step(message)

    def drop(self, from_: int, to: int, perc: float) -> None:
        """
        # Ignore messages from `from` to `to` at `perc` percent chance.
        #
        # `perc` set to `1f64` is a 100% chance, `0f64` is a 0% chance.
        """
        self.dropm[Connection(from_, to)] = perc

    def cut(self, one: int, other: int) -> None:
        """
        Cut the communication between the two given nodes.
        """
        self.drop(one, other, 1.0)
        self.drop(other, one, 1.0)

    def isolate(self, id: int) -> None:
        """
        Isolate the given raft to and from all other raft in the cluster.
        """
        for i in range(0, len(self.peers)):
            nid = i + 1

            if nid != id:
                self.drop(id, nid, 1.0)
                self.drop(nid, id, 1.0)

    def recover(self) -> None:
        """
        Recover the cluster conditions applied with `drop` and `ignore`.
        """
        self.dropm = {}
        self.ignorem = {}
