package merkle

import (
	"bytes"
	"crypto/sha256"
	"encoding/hex"
)

type Sibling struct {
	Hash string `json:"hash"`
	Left bool   `json:"left"`
}

func HashLeaf(leaf []byte) []byte {
	s := sha256.Sum256(leaf)

	return s[:]
}

func Root(leaves [][]byte) []byte {
	if len(leaves) == 0 {
		empty := sha256.Sum256(nil)

		return empty[:]
	}

	level := cloneLevel(leaves)

	for len(level) > 1 {
		next := make([][]byte, 0, (len(level)+1)/2)

		for i := 0; i < len(level); i += 2 {
			left := level[i]
			right := left

			if i+1 < len(level) {
				right = level[i+1]
			}

			next = append(next, hashPair(left, right))
		}

		level = next
	}

	return level[0]
}

func BuildProof(leaves [][]byte, index int) []Sibling {
	if len(leaves) == 0 || index < 0 || index >= len(leaves) {
		return nil
	}

	proof := make([]Sibling, 0)
	level := cloneLevel(leaves)
	idx := index

	for len(level) > 1 {
		if idx%2 == 0 {
			sib := idx + 1

			if sib >= len(level) {
				sib = idx
			}

			proof = append(proof, Sibling{Hash: hex.EncodeToString(level[sib]), Left: false})
		} else {
			sib := idx - 1
			proof = append(proof, Sibling{Hash: hex.EncodeToString(level[sib]), Left: true})
		}

		next := make([][]byte, 0, (len(level)+1)/2)

		for i := 0; i < len(level); i += 2 {
			left := level[i]
			right := left

			if i+1 < len(level) {
				right = level[i+1]
			}

			next = append(next, hashPair(left, right))
		}

		idx /= 2
		level = next
	}

	return proof
}

func VerifyProof(leafHash []byte, proof []Sibling, root []byte) bool {
	cur := append([]byte(nil), leafHash...)

	for _, s := range proof {
		sib, err := hex.DecodeString(s.Hash)

		if err != nil {
			return false
		}

		var combined []byte

		if s.Left {
			combined = append(append([]byte(nil), sib...), cur...)
		} else {
			combined = append(append([]byte(nil), cur...), sib...)
		}

		h := sha256.Sum256(combined)
		cur = h[:]
	}

	return bytes.Equal(cur, root)
}

func cloneLevel(in [][]byte) [][]byte {
	out := make([][]byte, len(in))

	for i := range in {
		out[i] = append([]byte(nil), in[i]...)
	}

	return out
}

type Accumulator struct {
	levels       [][]byte
	pendingCount int
}

func NewAccumulator() *Accumulator {
	return &Accumulator{levels: make([][]byte, 0, 32)}
}

func (a *Accumulator) AddLeafHash(leafHash []byte) {
	if len(leafHash) == 0 {
		return
	}

	cur := append([]byte(nil), leafHash...)
	level := 0

	for {
		if level >= len(a.levels) {
			a.levels = append(a.levels, nil)
		}

		if a.levels[level] == nil {
			a.levels[level] = cur

			a.pendingCount++

			break
		}

		cur = hashPair(a.levels[level], cur)
		a.levels[level] = nil

		a.pendingCount--
		level++
	}
}

func (a *Accumulator) RootDuplicateLast() []byte {
	if a.pendingCount == 0 {
		empty := sha256.Sum256(nil)

		return empty[:]
	}

	clone := &Accumulator{
		levels:       cloneLevel(a.levels),
		pendingCount: a.pendingCount,
	}

	for clone.pendingCount > 1 {
		level := clone.lowestPendingLevel()
		cur := clone.levels[level]
		clone.levels[level] = nil

		clone.pendingCount--

		cur = hashPair(cur, cur)

		level++

		for {
			if level >= len(clone.levels) {
				clone.levels = append(clone.levels, nil)
			}

			if clone.levels[level] == nil {
				clone.levels[level] = cur

				clone.pendingCount++

				break
			}

			cur = hashPair(clone.levels[level], cur)
			clone.levels[level] = nil

			clone.pendingCount--
			level++
		}
	}

	return clone.highestPendingHash()
}

func hashPair(left, right []byte) []byte {
	h := sha256.Sum256(append(append([]byte(nil), left...), right...))

	return h[:]
}

func (a *Accumulator) lowestPendingLevel() int {
	for i := 0; i < len(a.levels); i++ {
		if a.levels[i] != nil {
			return i
		}
	}

	return 0
}

func (a *Accumulator) highestPendingHash() []byte {
	for i := len(a.levels) - 1; i >= 0; i-- {
		if a.levels[i] != nil {
			return a.levels[i]
		}
	}

	empty := sha256.Sum256(nil)

	return empty[:]
}