package mls

import (
	"reflect"
	"testing"

	"github.com/cisco/go-tls-syntax"
	"github.com/stretchr/testify/require"
)

// Precomputed answers for the tree on eleven elements:
//
//                                              X
//                      X
//          X                       X                       X
//    X           X           X           X           X
// X     X     X     X     X     X     X     X     X     X     X
// 0  1  2  3  4  5  6  7  8  9  a  b  c  d  e  f 10 11 12 13 14
var (
	aRoot = []NodeIndex{0x00, 0x01, 0x03, 0x03, 0x07, 0x07, 0x07, 0x07, 0x0f, 0x0f, 0x0f}

	aN       = LeafCount(0x0b)
	index    = []NodeIndex{0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f, 0x10, 0x11, 0x12, 0x13, 0x14}
	aLog2    = []NodeIndex{0x00, 0x00, 0x01, 0x01, 0x02, 0x02, 0x02, 0x02, 0x03, 0x03, 0x03, 0x03, 0x03, 0x03, 0x03, 0x03, 0x04, 0x04, 0x04, 0x04, 0x04}
	aLevel   = []NodeIndex{0x00, 0x01, 0x00, 0x02, 0x00, 0x01, 0x00, 0x03, 0x00, 0x01, 0x00, 0x02, 0x00, 0x01, 0x00, 0x04, 0x00, 0x01, 0x00, 0x02, 0x00}
	aLeft    = []NodeIndex{0x00, 0x00, 0x02, 0x01, 0x04, 0x04, 0x06, 0x03, 0x08, 0x08, 0x0a, 0x09, 0x0c, 0x0c, 0x0e, 0x07, 0x10, 0x10, 0x12, 0x11, 0x14}
	aRight   = []NodeIndex{0x00, 0x02, 0x02, 0x05, 0x04, 0x06, 0x06, 0x0b, 0x08, 0x0a, 0x0a, 0x0d, 0x0c, 0x0e, 0x0e, 0x13, 0x10, 0x12, 0x12, 0x14, 0x14}
	aParent  = []NodeIndex{0x01, 0x03, 0x01, 0x07, 0x05, 0x03, 0x05, 0x0f, 0x09, 0x0b, 0x09, 0x07, 0x0d, 0x0b, 0x0d, 0x0f, 0x11, 0x13, 0x11, 0x0f, 0x13}
	aSibling = []NodeIndex{0x02, 0x05, 0x00, 0x0b, 0x06, 0x01, 0x04, 0x13, 0x0a, 0x0d, 0x08, 0x03, 0x0e, 0x09, 0x0c, 0x0f, 0x12, 0x14, 0x10, 0x07, 0x11}

	aDirpath = [][]NodeIndex{
		{0x01, 0x03, 0x07, 0x0f},
		{0x03, 0x07, 0x0f},
		{0x01, 0x03, 0x07, 0x0f},
		{0x07, 0x0f},
		{0x05, 0x03, 0x07, 0x0f},
		{0x03, 0x07, 0x0f},
		{0x05, 0x03, 0x07, 0x0f},
		{0x0f},
		{0x09, 0x0b, 0x07, 0x0f},
		{0x0b, 0x07, 0x0f},
		{0x09, 0x0b, 0x07, 0x0f},
		{0x07, 0x0f},
		{0x0d, 0x0b, 0x07, 0x0f},
		{0x0b, 0x07, 0x0f},
		{0x0d, 0x0b, 0x07, 0x0f},
		{},
		{0x11, 0x13, 0x0f},
		{0x13, 0x0f},
		{0x11, 0x13, 0x0f},
		{0x0f},
		{0x13, 0x0f},
	}
	aCopath = [][]NodeIndex{
		{0x02, 0x05, 0x0b, 0x13},
		{0x05, 0x0b, 0x13},
		{0x00, 0x05, 0x0b, 0x13},
		{0x0b, 0x13},
		{0x06, 0x01, 0x0b, 0x13},
		{0x01, 0x0b, 0x13},
		{0x04, 0x01, 0x0b, 0x13},
		{0x13},
		{0x0a, 0x0d, 0x03, 0x13},
		{0x0d, 0x03, 0x13},
		{0x08, 0x0d, 0x03, 0x13},
		{0x03, 0x13},
		{0x0e, 0x09, 0x03, 0x13},
		{0x09, 0x03, 0x13},
		{0x0c, 0x09, 0x03, 0x13},
		{},
		{0x12, 0x14, 0x07},
		{0x14, 0x07},
		{0x10, 0x14, 0x07},
		{0x07},
		{0x11, 0x07},
	}

	aInPath = [][]int{
		//    0  1  2  3  4  5  6  7  8  9  a  b  c  d  e  f 10 11 12 13 14
		/**/ {1, 1, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0}, //  0
		/**/ {0, 1, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0}, //  1
		/**/ {0, 1, 1, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0}, //  2
		/**/ {0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0}, //  3
		/**/ {0, 0, 0, 1, 1, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0}, //  4
		/**/ {0, 0, 0, 1, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0}, //  5
		/**/ {0, 0, 0, 1, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0}, //  6
		/**/ {0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0}, //  7
		/**/ {0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0}, //  8
		/**/ {0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0}, //  9
		/**/ {0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 1, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0}, //  a
		/**/ {0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0}, //  b
		/**/ {0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 1, 0, 1, 0, 0, 0, 0, 0}, //  c
		/**/ {0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 1, 0, 1, 0, 0, 0, 0, 0}, //  d
		/**/ {0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 1, 1, 1, 0, 0, 0, 0, 0}, //  e
		/**/ {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0}, //  f
		/**/ {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 1, 0}, // 10
		/**/ {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, 1, 0}, // 11
		/**/ {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 1, 1, 0}, // 12
		/**/ {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0}, // 13
		/**/ {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 1}, // 14
	}

	aFullAncestor = [][]NodeIndex{
		//  0     1     2     3     4     5     6     7     8     9     a     b     c     d     e     f    10    11    12    13    14
		{0x00, 0x01, 0x01, 0x03, 0x03, 0x03, 0x03, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f}, //  0
		{0x01, 0x01, 0x01, 0x03, 0x03, 0x03, 0x03, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f}, //  1
		{0x01, 0x01, 0x02, 0x03, 0x03, 0x03, 0x03, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f}, //  2
		{0x03, 0x03, 0x03, 0x03, 0x03, 0x03, 0x03, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f}, //  3
		{0x03, 0x03, 0x03, 0x03, 0x04, 0x05, 0x05, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f}, //  4
		{0x03, 0x03, 0x03, 0x03, 0x05, 0x05, 0x05, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f}, //  5
		{0x03, 0x03, 0x03, 0x03, 0x05, 0x05, 0x06, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f}, //  6
		{0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f}, //  7
		{0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x08, 0x09, 0x09, 0x0b, 0x0b, 0x0b, 0x0b, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f}, //  8
		{0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x09, 0x09, 0x09, 0x0b, 0x0b, 0x0b, 0x0b, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f}, //  9
		{0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x09, 0x09, 0x0a, 0x0b, 0x0b, 0x0b, 0x0b, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f}, //  a
		{0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x0b, 0x0b, 0x0b, 0x0b, 0x0b, 0x0b, 0x0b, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f}, //  b
		{0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x0b, 0x0b, 0x0b, 0x0b, 0x0c, 0x0d, 0x0d, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f}, //  c
		{0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x0b, 0x0b, 0x0b, 0x0b, 0x0d, 0x0d, 0x0d, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f}, //  d
		{0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x0b, 0x0b, 0x0b, 0x0b, 0x0d, 0x0d, 0x0e, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f}, //  e
		{0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f}, //  f
		{0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x10, 0x11, 0x11, 0x13, 0x13}, // 10
		{0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x11, 0x11, 0x11, 0x13, 0x13}, // 11
		{0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x11, 0x11, 0x12, 0x13, 0x13}, // 12
		{0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x13, 0x13, 0x13, 0x13, 0x13}, // 13
		{0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x0f, 0x13, 0x13, 0x13, 0x13, 0x14}, // 14
	}

	aAncestor = [][]NodeIndex{
		{0x01, 0x03, 0x03, 0x07, 0x07, 0x07, 0x07, 0x0f, 0x0f, 0x0f},
		{0x03, 0x03, 0x07, 0x07, 0x07, 0x07, 0x0f, 0x0f, 0x0f},
		{0x05, 0x07, 0x07, 0x07, 0x07, 0x0f, 0x0f, 0x0f},
		{0x07, 0x07, 0x07, 0x07, 0x0f, 0x0f, 0x0f},
		{0x09, 0x0b, 0x0b, 0x0f, 0x0f, 0x0f},
		{0x0b, 0x0b, 0x0f, 0x0f, 0x0f},
		{0x0d, 0x0f, 0x0f, 0x0f},
		{0x0f, 0x0f, 0x0f},
		{0x11, 0x13},
		{0x13},
	}
)

func TestSizeProperties(t *testing.T) {
	for n := LeafCount(1); n < aN; n += 1 {
		if root(n) != aRoot[n-1] {
			t.Fatalf("Root mismatch: %v != %v", root(n), aRoot[n-1])
		}
	}
}

func TestNodeRelations(t *testing.T) {
	run := func(label string, f func(x NodeIndex) NodeIndex, a []NodeIndex) {
		for i, x := range index {
			if f(x) != a[i] {
				t.Fatalf("Relation test failure: %s @ 0x%02x: %v != %v", label, x, f(x), a[i])
			}
		}
	}

	run("log2", func(x NodeIndex) NodeIndex { return NodeIndex(log2(nodeCount(x))) }, aLog2)
	run("level", func(x NodeIndex) NodeIndex { return NodeIndex(level(x)) }, aLevel)
	run("left", left, aLeft)
	run("right", func(x NodeIndex) NodeIndex { return right(x, aN) }, aRight)
	run("parent", func(x NodeIndex) NodeIndex { return parent(x, aN) }, aParent)
	run("sibling", func(x NodeIndex) NodeIndex { return sibling(x, aN) }, aSibling)
}

func TestPaths(t *testing.T) {
	run := func(label string, f func(x NodeIndex, n LeafCount) []NodeIndex, a [][]NodeIndex) {
		for i, x := range index {
			if !reflect.DeepEqual(f(x, aN), a[i]) {
				t.Fatalf("Path test failure: %s @ 0x%02x: %v != %v", label, x, f(x, aN), a[i])
			}
		}
	}

	run("dirpath", dirpath, aDirpath)
	run("copath", copath, aCopath)
}

func TestInPath(t *testing.T) {
	w := NodeIndex(nodeWidth(aN))
	for l := NodeIndex(0); l < w; l += 1 {
		for r := NodeIndex(0); r < w; r += 1 {
			answer := aInPath[l][r] == 1
			lr := inPath(l, r)

			if lr != answer {
				t.Errorf("Incorrect inPath determination: %d %d => %v != %v [%08b %08b]", l, r, lr, answer, l, r)
			}
		}
	}
}

func TestFullAncestor(t *testing.T) {
	w := NodeIndex(nodeWidth(aN))
	for l := NodeIndex(0); l < w; l += 1 {
		for r := NodeIndex(0); r < w; r += 1 {
			answer := aFullAncestor[l][r]
			lr := fullAncestor(l, r)
			rl := fullAncestor(r, l)

			if lr != answer {
				t.Errorf("Incorrect ancestor: %d %d => %d != %d", l, r, lr, answer)
			}

			if rl != lr {
				t.Errorf("Asymmetric ancestor: %d %d => %d != %d", l, r, rl, lr)
			}
		}
	}
}

func TestAncestor(t *testing.T) {
	for l := LeafIndex(0); l < LeafIndex(aN-1); l += 1 {
		for r := l + 1; r < LeafIndex(aN); r += 1 {
			answer := aAncestor[l][r-l-1]
			lr := ancestor(l, r)
			rl := ancestor(r, l)

			if lr != answer {
				t.Fatalf("Incorrect ancestor: %d %d => %d != %d", l, r, lr, answer)
			}

			if rl != answer {
				t.Fatalf("Asymmetric ancestor: %d %d => %d != %d", l, r, rl, lr)
			}
		}
	}
}

///
/// Test Vectors
///

type NodeIndexSlice struct {
	Data []NodeIndex `tls:"head=4"`
}

type TreeMathTestVectors struct {
	NumLeaves LeafCount
	Root      []NodeIndex      `tls:"head=4"`
	Left      []NodeIndex      `tls:"head=4"`
	Right     []NodeIndex      `tls:"head=4"`
	Parent    []NodeIndex      `tls:"head=4"`
	Sibling   []NodeIndex      `tls:"head=4"`
	DirPath   []NodeIndexSlice `tls:"head=4"`
	CoPath    []NodeIndexSlice `tls:"head=4"`
	Ancestor  []NodeIndexSlice `tls:"head=4"`
}

func generateTreeMathVectors(t *testing.T) []byte {
	numLeaves := LeafCount(255)
	numNodes := nodeWidth(numLeaves)
	tv := TreeMathTestVectors{
		NumLeaves: numLeaves,
		Root:      make([]NodeIndex, numLeaves),
		Left:      make([]NodeIndex, numNodes),
		Right:     make([]NodeIndex, numNodes),
		Parent:    make([]NodeIndex, numNodes),
		Sibling:   make([]NodeIndex, numNodes),
		DirPath:   make([]NodeIndexSlice, numNodes),
		CoPath:    make([]NodeIndexSlice, numNodes),
		Ancestor:  make([]NodeIndexSlice, numNodes),
	}

	for i := range tv.Root {
		tv.Root[i] = root(LeafCount(i + 1))
	}

	for i := range tv.Left {
		tv.Left[i] = left(NodeIndex(i))
		tv.Right[i] = right(NodeIndex(i), numLeaves)
		tv.Parent[i] = parent(NodeIndex(i), numLeaves)
		tv.Sibling[i] = sibling(NodeIndex(i), numLeaves)
		tv.DirPath[i] = NodeIndexSlice{Data: dirpath(NodeIndex(i), numLeaves)}
		tv.CoPath[i] = NodeIndexSlice{Data: copath(NodeIndex(i), numLeaves)}
	}

	// ancestor
	for l := LeafIndex(0); l < LeafIndex(numLeaves-1); l += 1 {
		a := []NodeIndex{}
		for r := l + 1; r < LeafIndex(numLeaves); r += 1 {
			lr := ancestor(l, r)
			a = append(a, lr)
		}
		tv.Ancestor[l].Data = a
	}

	vec, err := syntax.Marshal(tv)
	require.Nil(t, err)
	return vec
}

func verifyTreeMathVectors(t *testing.T, data []byte) {
	var tv TreeMathTestVectors
	_, err := syntax.Unmarshal(data, &tv)
	require.Nil(t, err)

	tvLen := int(nodeWidth(tv.NumLeaves))
	if len(tv.Root) != int(tv.NumLeaves) || len(tv.Left) != tvLen ||
		len(tv.Right) != tvLen || len(tv.Parent) != tvLen || len(tv.Sibling) != tvLen ||
		len(tv.DirPath) != tvLen || len(tv.CoPath) != tvLen {
		t.Fatalf("Malformed tree math test vectors: Incorrect vector sizes")
	}

	for i := range tv.Root {
		require.Equal(t, tv.Root[i], root(LeafCount(i+1)))
	}

	for i := range tv.Left {
		require.Equal(t, tv.Left[i], left(NodeIndex(i)))
		require.Equal(t, tv.Right[i], right(NodeIndex(i), tv.NumLeaves))
		require.Equal(t, tv.Parent[i], parent(NodeIndex(i), tv.NumLeaves))
		require.Equal(t, tv.Sibling[i], sibling(NodeIndex(i), tv.NumLeaves))
		require.Equal(t, tv.DirPath[i].Data, dirpath(NodeIndex(i), tv.NumLeaves))
		require.Equal(t, tv.CoPath[i].Data, copath(NodeIndex(i), tv.NumLeaves))
	}

	// ancestor
	for l := LeafIndex(0); l < LeafIndex(tv.NumLeaves-1); l += 1 {
		a := []NodeIndex{}
		for r := l + 1; r < LeafIndex(tv.NumLeaves); r += 1 {
			lr := ancestor(l, r)
			a = append(a, lr)
		}
		require.Equal(t, tv.Ancestor[l].Data, a)
	}

}

func TestTreeMathErrorCases(t *testing.T) {
	require.Panics(t, func() { toLeafIndex(0x03) })
}
