package wallet_test

import (
	"runtime"
	"testing"
	
	"github.com/p9c/p9/cmd/wallet"
)

// Harness holds the BranchRecoveryState being tested, the recovery window being used, provides access to the test
// object, and tracks the expected horizon and next unfound values.
type Harness struct {
	t              *testing.T
	brs            *wallet.BranchRecoveryState
	recoveryWindow uint32
	expHorizon     uint32
	expNextUnfound uint32
}
type (
	// Stepper is a generic interface that performs an action or assertion against a test Harness.
	Stepper interface {
		// Apply performs an action or assertion against branch recovery state held by the Harness. The step index is
		// provided so that any failures can report which Step failed.
		Apply(step int, harness *Harness)
	}
	// InitialiDelta is a Step that verifies our first attempt to expand the branch recovery state's horizons tells us
	// to derive a number of adddresses equal to the recovery window.
	InitialDelta struct{}
	// CheckDelta is a Step that expands the branch recovery state's horizon, and checks that the returned delta meets
	// our expected `delta`.
	CheckDelta struct {
		delta uint32
	}
	// CheckNumInvalid is a Step that asserts that the branch recovery state reports `total` invalid children with the
	// current horizon.
	CheckNumInvalid struct {
		total uint32
	}
	// MarkInvalid is a Step that marks the `child` as invalid in the branch recovery state.
	MarkInvalid struct {
		child uint32
	}
	// ReportFound is a Step that reports `child` as being found to the branch recovery state.
	ReportFound struct {
		child uint32
	}
)

// Apply extends the current horizon of the branch recovery state, and checks that the returned delta is equal to the
// test's recovery window. If the assertions pass, the harness's expected horizon is increased by the returned delta.
//
// NOTE: This should be used before applying any CheckDelta steps.
func (_ InitialDelta) Apply(i int, h *Harness) {
	curHorizon, delta := h.brs.ExtendHorizon()
	assertHorizon(h.t, i, curHorizon, h.expHorizon)
	assertDelta(h.t, i, delta, h.recoveryWindow)
	h.expHorizon += delta
}

// Apply extends the current horizon of the branch recovery state, and checks that the returned delta is equal to the
// CheckDelta's child value.
func (d CheckDelta) Apply(i int, h *Harness) {
	curHorizon, delta := h.brs.ExtendHorizon()
	assertHorizon(h.t, i, curHorizon, h.expHorizon)
	assertDelta(h.t, i, delta, d.delta)
	h.expHorizon += delta
}

// Apply queries the branch recovery state for the number of invalid children that lie between the last found address
// and the current horizon, and compares that to the CheckNumInvalid's total.
func (m CheckNumInvalid) Apply(i int, h *Harness) {
	assertNumInvalid(h.t, i, h.brs.NumInvalidInHorizon(), m.total)
}

// Apply marks the MarkInvalid's child index as invalid in the branch recovery state, and increments the harness's
// expected horizon.
func (m MarkInvalid) Apply(i int, h *Harness) {
	h.brs.MarkInvalidChild(m.child)
	h.expHorizon++
}

// Apply reports the ReportFound's child index as found in the branch recovery state. If the child index meets or
// exceeds our expected next unfound value, the expected value will be modified to be the child index + 1. Afterwards,
// this step asserts that the branch recovery state's next reported unfound value matches our potentially-updated value.
func (r ReportFound) Apply(i int, h *Harness) {
	h.brs.ReportFound(r.child)
	if r.child >= h.expNextUnfound {
		h.expNextUnfound = r.child + 1
	}
	assertNextUnfound(h.t, i, h.brs.NextUnfound(), h.expNextUnfound)
}

// Compile-time checks to ensure our steps implement the Step interface.
var _ Stepper = InitialDelta{}
var _ Stepper = CheckDelta{}
var _ Stepper = CheckNumInvalid{}
var _ Stepper = MarkInvalid{}
var _ Stepper = ReportFound{}

// TestBranchRecoveryState walks the BranchRecoveryState through a sequence of steps, verifying that:
//
//   - the horizon is properly expanded in response to found addrs
//
//   - report found children below or equal to previously found causes no change
//
//   - marking invalid children expands the horizon
func TestBranchRecoveryState(t *testing.T) {
	const recoveryWindow = 10
	recoverySteps := []Stepper{
		// First, check that expanding our horizon returns exactly the recovery window (10).
		InitialDelta{},
		// Expected horizon: 10. Report finding the 2nd addr, this should cause our horizon to expand by 2.
		ReportFound{1},
		CheckDelta{2},
		// Expected horizon: 12. Sanity check that expanding again reports zero delta, as nothing has changed.
		CheckDelta{0},
		// Now, report finding the 6th addr, which should expand our horizon to 16 with a detla of 4.
		ReportFound{5},
		CheckDelta{4},
		// Expected horizon: 16. Sanity check that expanding again reports zero delta, as nothing has changed.
		CheckDelta{0},
		// Report finding child index 5 again, nothing should change.
		ReportFound{5},
		CheckDelta{0},
		// Report finding a lower index that what was last found, nothing should change.
		ReportFound{4},
		CheckDelta{0},
		// Moving on, report finding the 11th addr, which should extend our horizon to 21.
		ReportFound{10},
		CheckDelta{5},
		// Expected horizon: 21. Before testing the lookahead expansion when encountering invalid child keys, check that
		// we are correctly starting with no invalid keys.
		CheckNumInvalid{0},
		// Now that the window has been expanded, simulate deriving invalid keys in range of addrs that are being
		// derived for the first time. The horizon will be incremented by one, as the recovery manager is expected to
		// try and derive at least the next address.
		MarkInvalid{17},
		CheckNumInvalid{1},
		CheckDelta{0},
		// Expected horizon: 22. Chk that deriving a second invalid key shows both invalid indexes currently within
		// the horizon.
		MarkInvalid{18},
		CheckNumInvalid{2},
		CheckDelta{0},
		// Expected horizon: 23. Lastly, report finding the addr immediately after our two invalid keys. This should
		// return our number of invalid keys within the horizon back to 0.
		ReportFound{19},
		CheckNumInvalid{0},
		// As the 20-th key was just marked found, our horizon will need to expand to 30. With the horizon at 23, the
		// delta returned should be 7.
		CheckDelta{7},
		CheckDelta{0},
		// Expected horizon: 30.
	}
	brs := wallet.NewBranchRecoveryState(recoveryWindow)
	harness := &Harness{
		t:              t,
		brs:            brs,
		recoveryWindow: recoveryWindow,
	}
	for i, step := range recoverySteps {
		step.Apply(i, harness)
	}
}
func assertHorizon(t *testing.T, i int, have, want uint32) {
	assertHaveWant(t, i, "incorrect horizon", have, want)
}
func assertDelta(t *testing.T, i int, have, want uint32) {
	assertHaveWant(t, i, "incorrect delta", have, want)
}
func assertNextUnfound(t *testing.T, i int, have, want uint32) {
	assertHaveWant(t, i, "incorrect next unfound", have, want)
}
func assertNumInvalid(t *testing.T, i int, have, want uint32) {
	assertHaveWant(t, i, "incorrect num invalid children", have, want)
}
func assertHaveWant(t *testing.T, i int, msg string, have, want uint32) {
	_, _, line, _ := runtime.Caller(2)
	if want != have {
		t.Fatalf(
			"[line: %d, step: %d] %s: got %d, want %d",
			line, i, msg, have, want,
		)
	}
}