package autopilot

import (
	"bytes"
	"io/ioutil"
	"os"
	"testing"
	"time"

	prand "math/rand"

	"github.com/btcsuite/btcd/btcec"
	"github.com/btcsuite/btcutil"
	"github.com/lightningnetwork/lnd/channeldb"
)

type genGraphFunc func() (testGraph, func(), error)

type testGraph interface {
	ChannelGraph

	addRandChannel(*btcec.PublicKey, *btcec.PublicKey,
		btcutil.Amount) (*ChannelEdge, *ChannelEdge, error)

	addRandNode() (*btcec.PublicKey, error)
}

func newDiskChanGraph() (testGraph, func(), error) {
	// First, create a temporary directory to be used for the duration of
	// this test.
	tempDirName, err := ioutil.TempDir("", "channeldb")
	if err != nil {
		return nil, nil, err
	}

	// Next, create channeldb for the first time.
	cdb, err := channeldb.Open(tempDirName)
	if err != nil {
		return nil, nil, err
	}

	cleanUp := func() {
		cdb.Close()
		os.RemoveAll(tempDirName)
	}

	return &databaseChannelGraph{
		db: cdb.ChannelGraph(),
	}, cleanUp, nil
}

var _ testGraph = (*databaseChannelGraph)(nil)

func newMemChanGraph() (testGraph, func(), error) {
	return newMemChannelGraph(), nil, nil
}

var _ testGraph = (*memChannelGraph)(nil)

var chanGraphs = []struct {
	name    string
	genFunc genGraphFunc
}{
	{
		name:    "disk_graph",
		genFunc: newDiskChanGraph,
	},
	{
		name:    "mem_graph",
		genFunc: newMemChanGraph,
	},
}

// TestPrefAttachmentSelectEmptyGraph ensures that when passed an
// empty graph, the NodeSores function always returns a score of 0.
func TestPrefAttachmentSelectEmptyGraph(t *testing.T) {
	const (
		maxChanSize = btcutil.Amount(btcutil.SatoshiPerBitcoin)
	)

	prefAttach := NewPrefAttachment()

	// Create a random public key, which we will query to get a score for.
	pub, err := randKey()
	if err != nil {
		t.Fatalf("unable to generate key: %v", err)
	}

	nodes := map[NodeID]struct{}{
		NewNodeID(pub): {},
	}

	for _, graph := range chanGraphs {
		success := t.Run(graph.name, func(t1 *testing.T) {
			graph, cleanup, err := graph.genFunc()
			if err != nil {
				t1.Fatalf("unable to create graph: %v", err)
			}
			if cleanup != nil {
				defer cleanup()
			}

			// With the necessary state initialized, we'll now
			// attempt to get the score for this one node.
			const walletFunds = btcutil.SatoshiPerBitcoin
			scores, err := prefAttach.NodeScores(graph, nil,
				walletFunds, nodes)
			if err != nil {
				t1.Fatalf("unable to select attachment "+
					"directives: %v", err)
			}

			// Since the graph is empty, we expect the score to be
			// 0, giving an empty return map.
			if len(scores) != 0 {
				t1.Fatalf("expected empty score map, "+
					"instead got %v ", len(scores))
			}
		})
		if !success {
			break
		}
	}
}

// completeGraph is a helper method that adds numNodes fully connected nodes to
// the graph.
func completeGraph(t *testing.T, g testGraph, numNodes int) {
	const chanCapacity = btcutil.SatoshiPerBitcoin
	nodes := make(map[int]*btcec.PublicKey)
	for i := 0; i < numNodes; i++ {
		for j := i + 1; j < numNodes; j++ {

			node1 := nodes[i]
			node2 := nodes[j]
			edge1, edge2, err := g.addRandChannel(
				node1, node2, chanCapacity)
			if err != nil {
				t.Fatalf("unable to generate channel: %v", err)
			}

			if node1 == nil {
				pubKeyBytes := edge1.Peer.PubKey()
				nodes[i], err = btcec.ParsePubKey(
					pubKeyBytes[:], btcec.S256(),
				)
				if err != nil {
					t.Fatalf("unable to parse pubkey: %v",
						err)
				}
			}

			if node2 == nil {
				pubKeyBytes := edge2.Peer.PubKey()
				nodes[j], err = btcec.ParsePubKey(
					pubKeyBytes[:], btcec.S256(),
				)
				if err != nil {
					t.Fatalf("unable to parse pubkey: %v",
						err)
				}
			}
		}
	}
}

// TestPrefAttachmentSelectTwoVertexes ensures that when passed a
// graph with only two eligible vertexes, then both are given the same score,
// and the funds are appropriately allocated across each peer.
func TestPrefAttachmentSelectTwoVertexes(t *testing.T) {
	t.Parallel()

	prand.Seed(time.Now().Unix())

	const (
		maxChanSize = btcutil.Amount(btcutil.SatoshiPerBitcoin)
	)

	for _, graph := range chanGraphs {
		success := t.Run(graph.name, func(t1 *testing.T) {
			graph, cleanup, err := graph.genFunc()
			if err != nil {
				t1.Fatalf("unable to create graph: %v", err)
			}
			if cleanup != nil {
				defer cleanup()
			}

			prefAttach := NewPrefAttachment()

			// For this set, we'll load the memory graph with two
			// nodes, and a random channel connecting them.
			const chanCapacity = btcutil.SatoshiPerBitcoin
			edge1, edge2, err := graph.addRandChannel(nil, nil, chanCapacity)
			if err != nil {
				t1.Fatalf("unable to generate channel: %v", err)
			}

			// We also add a third, non-connected node to the graph.
			_, err = graph.addRandNode()
			if err != nil {
				t1.Fatalf("unable to add random node: %v", err)
			}

			// Get the score for all nodes found in the graph at
			// this point.
			nodes := make(map[NodeID]struct{})
			if err := graph.ForEachNode(func(n Node) error {
				nodes[n.PubKey()] = struct{}{}
				return nil
			}); err != nil {
				t1.Fatalf("unable to traverse graph: %v", err)
			}

			if len(nodes) != 3 {
				t1.Fatalf("expected 2 nodes, found %d", len(nodes))
			}

			// With the necessary state initialized, we'll now
			// attempt to get our candidates channel score given
			// the current state of the graph.
			candidates, err := prefAttach.NodeScores(graph, nil,
				maxChanSize, nodes)
			if err != nil {
				t1.Fatalf("unable to select attachment "+
					"directives: %v", err)
			}

			// We expect two candidates, since one of the nodes
			// doesn't have any channels.
			if len(candidates) != 2 {
				t1.Fatalf("2 nodes should be scored, "+
					"instead %v were", len(candidates))
			}

			// The candidates should be amongst the two edges
			// created above.
			for nodeID, candidate := range candidates {
				edge1Pub := edge1.Peer.PubKey()
				edge2Pub := edge2.Peer.PubKey()

				switch {
				case bytes.Equal(nodeID[:], edge1Pub[:]):
				case bytes.Equal(nodeID[:], edge2Pub[:]):
				default:
					t1.Fatalf("attached to unknown node: %x",
						nodeID[:])
				}

				// Since each of the nodes has 1 channel, out
				// of only one channel in the graph, we expect
				// their score to be 1.0.
				expScore := float64(1.0)
				if candidate.Score != expScore {
					t1.Fatalf("expected candidate score "+
						"to be %v, instead was %v",
						expScore, candidate.Score)
				}
			}
		})
		if !success {
			break
		}
	}
}

// TestPrefAttachmentSelectGreedyAllocation tests that if upon
// returning node scores, the NodeScores method will attempt to greedily
// allocate all funds to each vertex (up to the max channel size).
func TestPrefAttachmentSelectGreedyAllocation(t *testing.T) {
	t.Parallel()

	prand.Seed(time.Now().Unix())

	const (
		maxChanSize = btcutil.Amount(btcutil.SatoshiPerBitcoin)
	)

	for _, graph := range chanGraphs {
		success := t.Run(graph.name, func(t1 *testing.T) {
			graph, cleanup, err := graph.genFunc()
			if err != nil {
				t1.Fatalf("unable to create graph: %v", err)
			}
			if cleanup != nil {
				defer cleanup()
			}

			prefAttach := NewPrefAttachment()

			const chanCapacity = btcutil.SatoshiPerBitcoin

			// Next, we'll add 3 nodes to the graph, creating an
			// "open triangle topology".
			edge1, _, err := graph.addRandChannel(nil, nil,
				chanCapacity)
			if err != nil {
				t1.Fatalf("unable to create channel: %v", err)
			}
			peerPubBytes := edge1.Peer.PubKey()
			peerPub, err := btcec.ParsePubKey(
				peerPubBytes[:], btcec.S256(),
			)
			if err != nil {
				t.Fatalf("unable to parse pubkey: %v", err)
			}
			_, _, err = graph.addRandChannel(
				peerPub, nil, chanCapacity,
			)
			if err != nil {
				t1.Fatalf("unable to create channel: %v", err)
			}

			// At this point, there should be three nodes in the
			// graph, with node node having two edges.
			numNodes := 0
			twoChans := false
			nodes := make(map[NodeID]struct{})
			if err := graph.ForEachNode(func(n Node) error {
				numNodes++
				nodes[n.PubKey()] = struct{}{}
				numChans := 0
				err := n.ForEachChannel(func(c ChannelEdge) error {
					numChans++
					return nil
				})
				if err != nil {
					return err
				}

				twoChans = twoChans || (numChans == 2)

				return nil
			}); err != nil {
				t1.Fatalf("unable to traverse graph: %v", err)
			}
			if numNodes != 3 {
				t1.Fatalf("expected 3 nodes, instead have: %v",
					numNodes)
			}
			if !twoChans {
				t1.Fatalf("expected node to have two channels")
			}

			// We'll now begin our test, modeling the available
			// wallet balance to be 5.5 BTC. We're shooting for a
			// 50/50 allocation, and have 3 BTC in channels. As a
			// result, the heuristic should try to greedily
			// allocate funds to channels.
			scores, err := prefAttach.NodeScores(graph, nil,
				maxChanSize, nodes)
			if err != nil {
				t1.Fatalf("unable to select attachment "+
					"directives: %v", err)
			}

			if len(scores) != len(nodes) {
				t1.Fatalf("all nodes should be scored, "+
					"instead %v were", len(scores))
			}

			// The candidates should have a non-zero score, and
			// have the max chan size funds recommended channel
			// size.
			for _, candidate := range scores {
				if candidate.Score == 0 {
					t1.Fatalf("Expected non-zero score")
				}
			}

			// Imagine a few channels are being opened, and there's
			// only 0.5 BTC left. That should leave us with channel
			// candidates of that size.
			const remBalance = btcutil.SatoshiPerBitcoin * 0.5
			scores, err = prefAttach.NodeScores(graph, nil,
				remBalance, nodes)
			if err != nil {
				t1.Fatalf("unable to select attachment "+
					"directives: %v", err)
			}

			if len(scores) != len(nodes) {
				t1.Fatalf("all nodes should be scored, "+
					"instead %v were", len(scores))
			}

			// Check that the recommended channel sizes are now the
			// remaining channel balance.
			for _, candidate := range scores {
				if candidate.Score == 0 {
					t1.Fatalf("Expected non-zero score")
				}
			}
		})
		if !success {
			break
		}
	}
}

// TestPrefAttachmentSelectSkipNodes ensures that if a node was
// already selected as a channel counterparty, then that node will get a score
// of zero during scoring.
func TestPrefAttachmentSelectSkipNodes(t *testing.T) {
	t.Parallel()

	prand.Seed(time.Now().Unix())

	const (
		maxChanSize = btcutil.Amount(btcutil.SatoshiPerBitcoin)
	)

	for _, graph := range chanGraphs {
		success := t.Run(graph.name, func(t1 *testing.T) {
			graph, cleanup, err := graph.genFunc()
			if err != nil {
				t1.Fatalf("unable to create graph: %v", err)
			}
			if cleanup != nil {
				defer cleanup()
			}

			prefAttach := NewPrefAttachment()

			// Next, we'll create a simple topology of two nodes,
			// with a single channel connecting them.
			const chanCapacity = btcutil.SatoshiPerBitcoin
			_, _, err = graph.addRandChannel(nil, nil,
				chanCapacity)
			if err != nil {
				t1.Fatalf("unable to create channel: %v", err)
			}

			nodes := make(map[NodeID]struct{})
			if err := graph.ForEachNode(func(n Node) error {
				nodes[n.PubKey()] = struct{}{}
				return nil
			}); err != nil {
				t1.Fatalf("unable to traverse graph: %v", err)
			}

			if len(nodes) != 2 {
				t1.Fatalf("expected 2 nodes, found %d", len(nodes))
			}

			// With our graph created, we'll now get the scores for
			// all nodes in the graph.
			scores, err := prefAttach.NodeScores(graph, nil,
				maxChanSize, nodes)
			if err != nil {
				t1.Fatalf("unable to select attachment "+
					"directives: %v", err)
			}

			if len(scores) != len(nodes) {
				t1.Fatalf("all nodes should be scored, "+
					"instead %v were", len(scores))
			}

			// THey should all have a score, and a maxChanSize
			// channel size recommendation.
			for _, candidate := range scores {
				if candidate.Score == 0 {
					t1.Fatalf("Expected non-zero score")
				}
			}

			// We'll simulate a channel update by adding the nodes
			// to our set of channels.
			var chans []Channel
			for _, candidate := range scores {
				chans = append(chans,
					Channel{
						Node: candidate.NodeID,
					},
				)
			}

			// If we attempt to make a call to the NodeScores
			// function, without providing any new information,
			// then all nodes should have a score of zero, since we
			// already got channels to them.
			scores, err = prefAttach.NodeScores(graph, chans,
				maxChanSize, nodes)
			if err != nil {
				t1.Fatalf("unable to select attachment "+
					"directives: %v", err)
			}

			// Since all should be given a score of 0, the map
			// should be empty.
			if len(scores) != 0 {
				t1.Fatalf("expected empty score map, "+
					"instead got %v ", len(scores))
			}
		})
		if !success {
			break
		}
	}
}