"""Integration tests for internal pathfinding."""
from typing import Dict
from unittest import TestCase, mock

import networkx as nx

from test import testing_common

from lndmanage.lib.data_types import NodePair
from lndmanage.lib.network import Network
from lndmanage.lib.rating import ChannelRater
from lndmanage.lib.pathfinding import dijkstra
from lndmanage.lib.liquidityhints import LiquidityHintMgr

from test.graph_definitions.routing_graph import nodes as test_graph


def new_test_graph(graph: Dict):
    # we need to init the node interface with a public key
    class MockNode:
        pub_key = 'A'
        query_mc = lambda x: {}

    # we disable cached graph reading
    with mock.patch.object(Network, 'load_graph', return_value=None):
        network = Network(MockNode())
        network.graph = nx.MultiGraph()
        network.edges = {}
        # TODO: fix side effects of liquidity hints loading
        network.liquidity_hints = LiquidityHintMgr(MockNode.pub_key)
        network.max_pair_capacity = {}

    # add nodes
    for node, node_definition in graph.items():
        network.graph.add_node(
            node,
            alias=node,
            last_update=None,
            address=None,
            color=None)

    # add channels
    for node, node_definition in graph.items():
        for channel, channel_definition in node_definition['channels'].items():
            # create a dictionary for channel_id lookups
            to_node = channel_definition['to']
            node_pair = NodePair((node, to_node))
            network.edges[channel] = {
                'node1_pub': node,
                'node2_pub': to_node,
                'node_pair': NodePair((node, to_node)),
                'capacity': channel_definition['capacity'],
                'last_update': None,
                'channel_id': channel,
                'chan_point': channel,
                'policies': {
                    node > to_node: channel_definition['policies'][node > to_node],
                    to_node > node: channel_definition['policies'][to_node > node],
                }
            }

            # add vertices to network graph for edge-based lookups
            network.graph.add_edge(
                node,
                to_node,
                channel_id=channel,
                last_update=None,
                capacity=channel_definition['capacity'],
                node_pair=NodePair((node, to_node)),
                fees={
                    node > to_node: channel_definition['policies'][node > to_node],
                    to_node > node: channel_definition['policies'][to_node > node],
                })

            # max channel capacity
            if not network.max_pair_capacity.get(node_pair):
                network.max_pair_capacity[node_pair] = channel_definition['capacity']
            else:
                if network.max_pair_capacity[node_pair] < channel_definition['capacity']:
                   network.max_pair_capacity[node_pair] = channel_definition['capacity']

    return network


class TestGraph(TestCase):
    def test_network(self):
        n = new_test_graph(test_graph)
        self.assertEqual(5, n.graph.number_of_nodes())
        self.assertEqual(7, n.graph.number_of_edges())

    def test_shortest_path(self):
        network = new_test_graph(test_graph)
        cr = ChannelRater(network)
        amt_msat = 1_000_000
        weight_function = lambda v, u, e: cr.node_to_node_weight(v, u, e, amt_msat)
        print(dijkstra(network.graph, 'A', 'E', weight=weight_function))
        # TODO: use too high capacity
        # TODO: use parallel channels with different policies

    def test_liquidity_hints(self):
        """
                3
            A  ---  B
            |    2/ |
          6 |   E   | 1
            | /5 \7 |
            D  ---  C
                4
        """
        amt_msat = 100_000 * 1_000
        network = new_test_graph(test_graph)
        cr = ChannelRater(network=network)
        weight_function = lambda v, u, e: cr.node_to_node_weight(v, u, e, amt_msat)

        path = dijkstra(network.graph, 'A', 'E', weight=weight_function)
        self.assertEqual(['A', 'B', 'E'], path)

        # We report that B cannot send to E
        network.liquidity_hints.update_cannot_send('B', 'E', 1_000)
        path = dijkstra(network.graph, 'A', 'E', weight=weight_function)
        self.assertEqual(['A', 'D', 'E'], path)

        # We report that D cannot send to E
        network.liquidity_hints.update_cannot_send('D', 'E', 1_000)
        path = dijkstra(network.graph, 'A', 'E', weight=weight_function)
        self.assertEqual(['A', 'B', 'C', 'E'], path)

        # We report that D can send to C
        network.liquidity_hints.update_can_send('D', 'C', amt_msat + 1000)
        path = dijkstra(network.graph, 'A', 'E', weight=weight_function)
        self.assertEqual(['A', 'D', 'C', 'E'], path)
