package itest

import (
	"encoding/hex"
	"fmt"
	"time"

	"github.com/btcsuite/btcd/btcutil"
	"github.com/lightningnetwork/lnd/chainreg"
	"github.com/lightningnetwork/lnd/lnrpc"
	"github.com/lightningnetwork/lnd/lnrpc/routerrpc"
	"github.com/lightningnetwork/lnd/lntest"
	"github.com/lightningnetwork/lnd/lntest/node"
	"github.com/lightningnetwork/lnd/lntest/wait"
	"github.com/lightningnetwork/lnd/lntypes"
	"github.com/lightningnetwork/lnd/routing/route"
	"github.com/stretchr/testify/require"
)

// testSendMultiPathPayment tests that we are able to successfully route a
// payment using multiple shards across different paths.
func testSendMultiPathPayment(ht *lntest.HarnessTest) {
	mts := newMppTestScenario(ht)

	// Set up a network with three different paths Alice <-> Bob. Channel
	// capacities are set such that the payment can only succeed if (at
	// least) three paths are used.
	//
	//              _ Eve _
	//             /       \
	// Alice -- Carol ---- Bob
	//      \              /
	//       \__ Dave ____/
	//
	paymentAmt := mts.setupSendPaymentCase()

	chanPointAliceDave := mts.channelPoints[1]

	// Increase Dave's fee to make the test deterministic. Otherwise, it
	// would be unpredictable whether pathfinding would go through Charlie
	// or Dave for the first shard.
	expectedPolicy := &lnrpc.RoutingPolicy{
		FeeBaseMsat:      500_000,
		FeeRateMilliMsat: int64(0.001 * 1_000_000),
		TimeLockDelta:    40,
		MinHtlc:          1000, // default value
		MaxHtlcMsat:      133_650_000,
	}
	mts.dave.UpdateGlobalPolicy(expectedPolicy)

	// Make sure Alice has heard it.
	ht.AssertChannelPolicyUpdate(
		mts.alice, mts.dave, expectedPolicy, chanPointAliceDave, false,
	)

	// Our first test will be Alice paying Bob using a SendPayment call.
	// Let Bob create an invoice for Alice to pay.
	payReqs, rHashes, invoices := ht.CreatePayReqs(mts.bob, paymentAmt, 1)

	rHash := rHashes[0]
	payReq := payReqs[0]

	sendReq := &routerrpc.SendPaymentRequest{
		PaymentRequest: payReq,
		MaxParts:       10,
		TimeoutSeconds: 60,
		FeeLimitMsat:   noFeeLimitMsat,
	}
	payment := ht.SendPaymentAssertSettled(mts.alice, sendReq)

	// Make sure we got the preimage.
	require.Equal(ht, hex.EncodeToString(invoices[0].RPreimage),
		payment.PaymentPreimage, "preimage doesn't match")

	// Check that Alice split the payment in at least three shards. Because
	// the hand-off of the htlc to the link is asynchronous (via a mailbox),
	// there is some non-determinism in the process. Depending on whether
	// the new pathfinding round is started before or after the htlc is
	// locked into the channel, different sharding may occur. Therefore we
	// can only check if the number of shards isn't below the theoretical
	// minimum.
	succeeded := 0
	for _, htlc := range payment.Htlcs {
		if htlc.Status == lnrpc.HTLCAttempt_SUCCEEDED {
			succeeded++
		}
	}

	const minExpectedShards = 3
	require.GreaterOrEqual(ht, succeeded, minExpectedShards,
		"expected shards not reached")

	// Make sure Bob show the invoice as settled for the full amount.
	inv := mts.bob.RPC.LookupInvoice(rHash)

	require.EqualValues(ht, paymentAmt, inv.AmtPaidSat,
		"incorrect payment amt")

	require.Equal(ht, lnrpc.Invoice_SETTLED, inv.State,
		"Invoice not settled")

	settled := 0
	for _, htlc := range inv.Htlcs {
		if htlc.State == lnrpc.InvoiceHTLCState_SETTLED {
			settled++
		}
	}
	require.Equal(ht, succeeded, settled,
		"num of HTLCs wrong")

	// Finally, close all channels.
	mts.closeChannels()
}

// testSendToRouteMultiPath tests that we are able to successfully route a
// payment using multiple shards across different paths, by using SendToRoute.
func testSendToRouteMultiPath(ht *lntest.HarnessTest) {
	mts := newMppTestScenario(ht)

	// To ensure the payment goes through separate paths, we'll set a
	// channel size that can only carry one shard at a time. We'll divide
	// the payment into 3 shards.

	// Set up a network with three different paths Alice <-> Bob.
	//              _ Eve _
	//             /       \
	// Alice -- Carol ---- Bob
	//      \              /
	//       \__ Dave ____/
	//
	paymentAmt, shardAmt := mts.setupSendToRouteCase()

	// Make Bob create an invoice for Alice to pay.
	payReqs, rHashes, invoices := ht.CreatePayReqs(mts.bob, paymentAmt, 1)

	rHash := rHashes[0]
	payReq := payReqs[0]

	decodeResp := mts.bob.RPC.DecodePayReq(payReq)
	payAddr := decodeResp.PaymentAddr

	// Subscribe the invoice.
	stream := mts.bob.RPC.SubscribeSingleInvoice(rHash)

	// We'll send shards along three routes from Alice.
	sendRoutes := [][]*node.HarnessNode{
		{mts.carol, mts.bob},
		{mts.dave, mts.bob},
		{mts.carol, mts.eve, mts.bob},
	}

	responses := make(chan *lnrpc.HTLCAttempt, len(sendRoutes))
	for _, hops := range sendRoutes {
		// Build a route for the specified hops.
		r := mts.buildRoute(shardAmt, mts.alice, hops)

		// Set the MPP records to indicate this is a payment shard.
		hop := r.Hops[len(r.Hops)-1]
		hop.TlvPayload = true
		hop.MppRecord = &lnrpc.MPPRecord{
			PaymentAddr:  payAddr,
			TotalAmtMsat: int64(paymentAmt * 1000),
		}

		// Send the shard.
		sendReq := &routerrpc.SendToRouteRequest{
			PaymentHash: rHash,
			Route:       r,
		}

		// We'll send all shards in their own goroutine, since
		// SendToRoute will block as long as the payment is in flight.
		go func() {
			resp := mts.alice.RPC.SendToRouteV2(sendReq)
			responses <- resp
		}()
	}

	// Wait for all responses to be back, and check that they all
	// succeeded.
	timer := time.After(defaultTimeout)
	for range sendRoutes {
		var resp *lnrpc.HTLCAttempt
		select {
		case resp = <-responses:
		case <-timer:
			require.Fail(ht, "response not received")
		}

		require.Nil(ht, resp.Failure, "received payment failure")

		// All shards should come back with the preimage.
		require.Equal(ht, resp.Preimage, invoices[0].RPreimage,
			"preimage doesn't match")
	}

	// assertNumHtlcs is a helper that checks the node's latest payment,
	// and asserts it was split into num shards.
	assertNumHtlcs := func(hn *node.HarnessNode, num int) {
		var preimage lntypes.Preimage
		copy(preimage[:], invoices[0].RPreimage)

		payment := ht.AssertPaymentStatus(
			hn, preimage, lnrpc.Payment_SUCCEEDED,
		)

		htlcs := payment.Htlcs
		require.NotEmpty(ht, htlcs, "no htlcs")

		succeeded := 0
		for _, htlc := range htlcs {
			if htlc.Status == lnrpc.HTLCAttempt_SUCCEEDED {
				succeeded++
			}
		}
		require.Equal(ht, num, succeeded, "HTLCs not matched")
	}

	// assertSettledInvoice checks that the invoice for the given payment
	// hash is settled, and has been paid using num HTLCs.
	assertSettledInvoice := func(rhash []byte, num int) {
		var payHash lntypes.Hash
		copy(payHash[:], rhash)
		inv := ht.AssertInvoiceState(stream, lnrpc.Invoice_SETTLED)

		// Assert that the amount paid to the invoice is correct.
		require.EqualValues(ht, paymentAmt, inv.AmtPaidSat,
			"incorrect payment amt")

		require.Len(ht, inv.Htlcs, num, "wrong num of HTLCs")
	}

	// Finally check that the payment shows up with three settled HTLCs in
	// Alice's list of payments...
	assertNumHtlcs(mts.alice, 3)

	// ...and in Bob's list of paid invoices.
	assertSettledInvoice(rHash, 3)

	// Finally, close all channels.
	mts.closeChannels()
}

// mppTestScenario defines a test scenario used for testing MPP-related tests.
// It has two standby nodes, alice and bob, and three new nodes, carol, dave,
// and eve.
type mppTestScenario struct {
	ht *lntest.HarnessTest

	alice, bob, carol, dave, eve *node.HarnessNode
	nodes                        []*node.HarnessNode

	// Keep a list of all our active channels.
	channelPoints []*lnrpc.ChannelPoint
}

// newMppTestScenario initializes a new mpp test scenario with five funded
// nodes and connects them to have the following topology,
//
//	            _ Eve _
//	           /       \
//	Alice -- Carol ---- Bob
//	    \              /
//	     \__ Dave ____/
//
// The scenario is setup in a way that when sending a payment from Alice to
// Bob, (at least) three routes must be tried to succeed.
func newMppTestScenario(ht *lntest.HarnessTest) *mppTestScenario {
	alice := ht.NewNodeWithCoins("Alice", nil)
	bob := ht.NewNodeWithCoins("Bob", []string{
		"--maxpendingchannels=2",
		"--accept-amp",
	})

	// Create a five-node context consisting of Alice, Bob and three new
	// nodes.
	carol := ht.NewNode("carol", []string{
		"--maxpendingchannels=2",
		"--accept-amp",
	})
	dave := ht.NewNode("dave", nil)
	eve := ht.NewNode("eve", nil)

	// Connect nodes to ensure propagation of channels.
	ht.EnsureConnected(alice, carol)
	ht.EnsureConnected(alice, dave)
	ht.EnsureConnected(carol, bob)
	ht.EnsureConnected(carol, eve)
	ht.EnsureConnected(dave, bob)
	ht.EnsureConnected(eve, bob)

	// Send coins to the nodes and mine 1 blocks to confirm them.
	for i := 0; i < 2; i++ {
		ht.FundCoinsUnconfirmed(btcutil.SatoshiPerBitcoin, carol)
		ht.FundCoinsUnconfirmed(btcutil.SatoshiPerBitcoin, dave)
		ht.FundCoinsUnconfirmed(btcutil.SatoshiPerBitcoin, eve)

		ht.MineBlocksAndAssertNumTxes(1, 3)
	}

	mts := &mppTestScenario{
		ht:    ht,
		alice: alice,
		bob:   bob,
		carol: carol,
		dave:  dave,
		eve:   eve,
		nodes: []*node.HarnessNode{alice, bob, carol, dave, eve},
	}

	return mts
}

// mppOpenChannelRequest defines the amounts used for each channel opening.
type mppOpenChannelRequest struct {
	// Channel Alice=>Carol.
	amtAliceCarol btcutil.Amount

	// Channel Alice=>Dave.
	amtAliceDave btcutil.Amount

	// Channel Carol=>Bob.
	amtCarolBob btcutil.Amount

	// Channel Carol=>Eve.
	amtCarolEve btcutil.Amount

	// Channel Dave=>Bob.
	amtDaveBob btcutil.Amount

	// Channel Eve=>Bob.
	amtEveBob btcutil.Amount
}

// setupSendPaymentCase opens channels between the nodes for testing the
// `SendPaymentV2` case, where a payment amount of 300,000 sats is used and it
// tests sending three attempts: the first has 150,000 sats, the rest two have
// 75,000 sats. It returns the payment amt.
func (c *mppTestScenario) setupSendPaymentCase() btcutil.Amount {
	// To ensure the payment goes through separate paths, we'll set a
	// channel size that can only carry one HTLC attempt at a time. We'll
	// divide the payment into 3 attempts.
	//
	// Set the payment amount to be 300,000 sats. When a route cannot be
	// found for a given payment amount, we will halven the amount and try
	// the pathfinding again, which means we need to see the following
	// three attempts to succeed:
	// 1. 1st attempt: 150,000 sats.
	// 2. 2nd attempt: 75,000 sats.
	// 3. 3rd attempt: 75,000 sats.
	paymentAmt := btcutil.Amount(300_000)

	// Prepare to open channels between the nodes. Given our expected
	// topology,
	//
	//		      _ Eve _
	//		     /       \
	//	 Alice -- Carol ---- Bob
	//		\             /
	//		 \__ Dave ___/
	//
	// There are three routes from Alice to Bob:
	// 1. Alice -> Carol -> Bob
	// 2. Alice -> Dave -> Bob
	// 3. Alice -> Carol -> Eve -> Bob
	// We now use hardcoded amounts so it's easier to reason about the
	// test.
	req := &mppOpenChannelRequest{
		amtAliceCarol: 285_000,
		amtAliceDave:  155_000,
		amtCarolBob:   200_000,
		amtCarolEve:   155_000,
		amtDaveBob:    155_000,
		amtEveBob:     155_000,
	}

	// Given the above setup, the only possible routes to send each of the
	// attempts are:
	// - 1st attempt(150,000 sats): Alice->Carol->Bob: 200,000 sats.
	// - 2nd attempt(75,000 sats): Alice->Dave->Bob: 155,000 sats.
	// - 3rd attempt(75,000 sats): Alice->Carol->Eve->Bob: 155,000 sats.
	//
	// There is a case where the payment will fail due to the channel
	// bandwidth not being updated in the graph, which has been seen many
	// times:
	// 1. the 1st attempt (150,000 sats) is sent via
	//    Alice->Carol->Eve->Bob, after which the capacity in Carol->Eve
	//    should decrease.
	// 2. the 2nd attempt (75,000 sats) is sent via Alice->Carol->Eve->Bob,
	//    which shouldn't happen because the capacity in Carol->Eve is
	//    depleted. However, since the HTLCs are sent in parallel, the 2nd
	//    attempt can be sent before the capacity is updated in the graph.
	// 3. if the 2nd attempt succeeds, the 1st attempt will fail and be
	//    split into two attempts, each holding 75,000 sats. At this point,
	//    we have three attempts to send, but only two routes are
	//    available, causing the payment to be failed.
	// 4. In addition, with recent fee buffer addition, the attempts will
	//    fail even earlier without being further split.
	//
	// To avoid this case, we now increase the channel capacity of the
	// route Carol->Eve->Bob and Carol->Bob such that even the above case
	// happened, we can still send the HTLCs.
	//
	// TODO(yy): we should properly fix this in the router. Atm we only
	// perform this hack to unblock the CI.
	req.amtCarolBob = 285_000
	req.amtEveBob = 285_000
	req.amtCarolEve = 285_000

	// Open the channels as described above.
	c.openChannels(req)

	return paymentAmt
}

// setupSendToRouteCase opens channels between the nodes for testing the
// `SendToRouteV2` case, where a payment amount of 300,000 sats is used and it
// tests sending three attempts each holding 100,000 sats. It returns the
// payment amount and attempt amount.
func (c *mppTestScenario) setupSendToRouteCase() (btcutil.Amount,
	btcutil.Amount) {

	// To ensure the payment goes through separate paths, we'll set a
	// channel size that can only carry one HTLC attempt at a time. We'll
	// divide the payment into 3 attempts, each holding 100,000 sats.
	paymentAmt := btcutil.Amount(300_000)
	attemptAmt := btcutil.Amount(100_000)

	// Prepare to open channels between the nodes. Given our expected
	// topology,
	//
	//		      _ Eve _
	//		     /       \
	//	 Alice -- Carol ---- Bob
	//		\             /
	//		 \__ Dave ___/
	//
	// There are three routes from Alice to Bob:
	// 1. Alice -> Carol -> Bob
	// 2. Alice -> Dave -> Bob
	// 3. Alice -> Carol -> Eve -> Bob
	// We now use hardcoded amounts so it's easier to reason about the
	// test.
	req := &mppOpenChannelRequest{
		amtAliceCarol: 250_000,
		amtAliceDave:  150_000,
		amtCarolBob:   150_000,
		amtCarolEve:   150_000,
		amtDaveBob:    150_000,
		amtEveBob:     150_000,
	}

	// Given the above setup, the only possible routes to send each of the
	// attempts are:
	// - 1st attempt(100,000 sats): Alice->Carol->Bob: 150,000 sats.
	// - 2nd attempt(100,000 sats): Alice->Dave->Bob: 150,000 sats.
	// - 3rd attempt(100,000 sats): Alice->Carol->Eve->Bob: 150,000 sats.
	//
	// Open the channels as described above.
	c.openChannels(req)

	return paymentAmt, attemptAmt
}

// openChannels is a helper to open channels that sets up a network topology
// with three different paths Alice <-> Bob as following,
//
//		      _ Eve _
//		     /       \
//	 Alice -- Carol ---- Bob
//		\              /
//		 \__ Dave ____/
//
// NOTE: all the channels are open together to save blocks mined.
func (m *mppTestScenario) openChannels(r *mppOpenChannelRequest) {
	reqs := []*lntest.OpenChannelRequest{
		{
			Local:  m.alice,
			Remote: m.carol,
			Param:  lntest.OpenChannelParams{Amt: r.amtAliceCarol},
		},
		{
			Local:  m.alice,
			Remote: m.dave,
			Param:  lntest.OpenChannelParams{Amt: r.amtAliceDave},
		},
		{
			Local:  m.carol,
			Remote: m.bob,
			Param:  lntest.OpenChannelParams{Amt: r.amtCarolBob},
		},
		{
			Local:  m.carol,
			Remote: m.eve,
			Param:  lntest.OpenChannelParams{Amt: r.amtCarolEve},
		},
		{
			Local:  m.dave,
			Remote: m.bob,
			Param:  lntest.OpenChannelParams{Amt: r.amtDaveBob},
		},
		{
			Local:  m.eve,
			Remote: m.bob,
			Param:  lntest.OpenChannelParams{Amt: r.amtEveBob},
		},
	}

	m.channelPoints = m.ht.OpenMultiChannelsAsync(reqs)

	// Make sure every node has heard every channel.
	for _, hn := range m.nodes {
		for _, cp := range m.channelPoints {
			m.ht.AssertChannelInGraph(hn, cp)
		}

		// Each node should have exactly 6 edges.
		m.ht.AssertNumEdges(hn, len(m.channelPoints), false)
	}
}

// closeChannels closes all the open channels from `openChannels`.
func (m *mppTestScenario) closeChannels() {
	if m.ht.Failed() {
		m.ht.Log("Skipped closing channels for failed test")
		return
	}

	// Close all channels without mining the closing transactions.
	m.ht.CloseChannelAssertPending(m.alice, m.channelPoints[0], false)
	m.ht.CloseChannelAssertPending(m.alice, m.channelPoints[1], false)
	m.ht.CloseChannelAssertPending(m.carol, m.channelPoints[2], false)
	m.ht.CloseChannelAssertPending(m.carol, m.channelPoints[3], false)
	m.ht.CloseChannelAssertPending(m.dave, m.channelPoints[4], false)
	m.ht.CloseChannelAssertPending(m.eve, m.channelPoints[5], false)

	// Now mine a block to include all the closing transactions.
	m.ht.MineBlocksAndAssertNumTxes(1, 6)

	// Assert that the channels are closed.
	for _, hn := range m.nodes {
		m.ht.AssertNumWaitingClose(hn, 0)
	}
}

// Helper function for Alice to build a route from pubkeys.
func (m *mppTestScenario) buildRoute(amt btcutil.Amount,
	sender *node.HarnessNode, hops []*node.HarnessNode) *lnrpc.Route {

	rpcHops := make([][]byte, 0, len(hops))
	for _, hop := range hops {
		k := hop.PubKeyStr
		pubkey, err := route.NewVertexFromStr(k)
		require.NoErrorf(m.ht, err, "error parsing %v: %v", k, err)
		rpcHops = append(rpcHops, pubkey[:])
	}

	req := &routerrpc.BuildRouteRequest{
		AmtMsat:        int64(amt * 1000),
		FinalCltvDelta: chainreg.DefaultBitcoinTimeLockDelta,
		HopPubkeys:     rpcHops,
	}

	// We should be able to call `sender.RPC.BuildRoute` directly, but
	// sometimes we will get a RPC-level error saying we cannot find the
	// node index:
	// - no matching outgoing channel available for node index 1
	// This happens because the `getEdgeUnifiers` cannot find a policy for
	// one of the hops,
	// - [ERR] CRTR router.go:1689: Cannot find policy for node ...
	// However, by the time we get here, we have already checked that all
	// nodes have heard all channels, so this indicates a bug in our
	// pathfinding, specifically in the edge unifier.
	//
	// TODO(yy): Remove the following wait and use the direct call, then
	// investigate the bug in the edge unifier.
	var route *lnrpc.Route
	err := wait.NoError(func() error {
		routeResp, err := sender.RPC.Router.BuildRoute(
			m.ht.Context(), req,
		)
		if err != nil {
			return fmt.Errorf("unable to build route for %v "+
				"using hops=%v: %v", sender.Name(), hops, err)
		}

		route = routeResp.Route

		return nil
	}, defaultTimeout)
	require.NoError(m.ht, err, "build route timeout")

	return route
}