package main

import (
	"bytes"
	"encoding/binary"
	"os"
	"sync"

	"golang.org/x/crypto/sha3"
)

var hashSize = int64(sha3.New256().Size())

type Merkle interface {
	Build(parallel bool) []byte
	Put(id int64, data []byte)
	Read(buf []byte, id int64)
	Size() int64
	Proof(id int64) [][]byte
}

func NewMerkle(mtype string, height int64, fname string) Merkle {
	var m Merkle
	switch mtype {
	case "bfs":
		m = NewBFSMerkle(height, fname)
	case "post":
		m = NewPostMerkle(height, fname)
	}
	return m
}

// nodes are stored in BFS order, root node first
type BFSMerkle struct {
	height int64 // root counts as height 1, children of root as height 2, etc.
	*os.File
	size int64 // maximum label of node = 2^height -2
}

func NewBFSMerkle(height int64, fname string) *BFSMerkle {
	file, err := os.OpenFile(fname, os.O_RDWR|os.O_CREATE, 0666)
	if err != nil {
		panic(err)
	}
	size := int64(1)<<uint64(height) - 2
	return &BFSMerkle{height: height, File: file, size: size}
}

func (m *BFSMerkle) Size() int64 {
	return m.size
}

func (m *BFSMerkle) Put(id int64, data []byte) {
	m.WriteAt(data, id*hashSize)
}

func (m *BFSMerkle) Read(buf []byte, id int64) {
	m.ReadAt(buf, id*hashSize)
}

func (m *BFSMerkle) Build(parallel bool) []byte {
	if parallel {
		return m.parallelbuild()
	} else {
		return m.build()
	}
}

// disk access is sequential and mostly backward
func (m *BFSMerkle) build() []byte {
	size := m.Size()
	h := sha3.New256()
	hsize := int64(h.Size())
	buf := make([]byte, hsize)

	for id := m.Size() / 2; id <= m.Size(); id++ {
		h.Reset()
		binary.Write(h, binary.LittleEndian, id)
		buf = h.Sum(buf[:0])
		m.WriteAt(buf, id*hsize)
	}

	for id := size/2 - 1; id >= 0; id-- {
		h.Reset()
		m.ReadAt(buf, (id*2+1)*hsize)
		h.Write(buf)
		m.ReadAt(buf, (id*2+2)*hsize)
		h.Write(buf)
		buf = h.Sum(buf[:0])
		m.WriteAt(buf, id*hsize)
	}
	return buf
}

func (m *BFSMerkle) parallelbuild() []byte {
	size := m.Size()
	hsize := hashSize
	var left int64
	var wg sync.WaitGroup
	p := 100

	i := 0
	for id := size / 2; id <= size; id++ {
		wg.Add(1)
		go func(id int64) {
			defer wg.Done()
			var buf bytes.Buffer
			binary.Write(&buf, binary.LittleEndian, id)
			hash := sha3.Sum256(buf.Bytes())
			m.WriteAt(hash[:], id*hsize)
		}(id)
		if i%p == 0 {
			wg.Wait()
		}
		i++
	}
	wg.Wait()

	for height := m.height - 1; height >= 1; height-- {
		left = 1 << uint64(height-1)
		i := 0
		for id := left - 1; id < 2*left-1; id++ {
			wg.Add(1)
			go func(id int64) {
				defer wg.Done()
				buf := make([]byte, 2*hsize)
				m.ReadAt(buf[:hsize], (id*2+1)*hsize)
				m.ReadAt(buf[hsize:], (id*2+2)*hsize)
				hash := sha3.Sum256(buf)
				m.WriteAt(hash[:], id*hsize)
			}(id)
			if i%p == 0 {
				wg.Wait()
			}
			i++
		}
		wg.Wait()
	}
	buf := make([]byte, hsize)
	m.ReadAt(buf, 0)
	return buf
}

func (m *BFSMerkle) Proof(id int64) [][]byte {
	proof := make([][]byte, m.height)
	proof[0] = make([]byte, hashSize)
	m.ReadAt(proof[0], id*hashSize)                       // read the queried node
	for height := int64(1); height < m.height; height++ { //construct the proof bottom-up
		proof[height] = make([]byte, hashSize)
		if id&1 == 0 { // right child, reading left sibling
			m.ReadAt(proof[height], (id-1)*hashSize)
		} else { // left child, reading right sibling
			m.ReadAt(proof[height], (id+1)*hashSize)
		}
		id = (id - 1) >> 1 // move to parent node
	}
	return proof
}

func (m *BFSMerkle) Proofs(ids []int64) [][][]byte {
	proofs := make([][][]byte, len(ids))
	for i, id := range ids {
		proofs[i] = m.Proof(id)
	}
	return proofs
}

func (m *BFSMerkle) BatchProofs(ids []int64) [][][]byte {
	proofs := make([][][]byte, len(ids))
	tids := make([]int64, len(ids))
	copy(tids, ids)
	var id int64

	// pre-allocating proofs + first step
	for i := range proofs {
		proofs[i] = make([][]byte, m.height)
		proofs[i][0] = make([]byte, hashSize)
		m.ReadAt(proofs[i][0], tids[i]*hashSize)
	}

	// constructing proofs level by level
	for height := int64(1); height < m.height; height++ {
		for i := range proofs {
			proofs[i][height] = make([]byte, hashSize)
			id = tids[i]
			if i > 0 && id == tids[i-1] {
				proofs[i][height] = proofs[i-1][height]
			} else if id&1 == 0 {
				m.ReadAt(proofs[i][height], (id-1)*hashSize)
			} else {
				m.ReadAt(proofs[i][height], (id+1)*hashSize)
			}
			tids[i] = (id - 1) >> 1
		}
	}
	return proofs
}

// nodes are stored in depth-first post order
type PostMerkle struct {
	height int64
	*os.File
	size int64
}

func NewPostMerkle(height int64, fname string) *PostMerkle {
	file, err := os.OpenFile(fname, os.O_RDWR|os.O_CREATE, 0666)
	if err != nil {
		panic(err)
	}
	size := int64(1)<<uint64(height) - 2
	return &PostMerkle{height: height, File: file, size: size}
}

func (m *PostMerkle) Size() int64 {
	return m.size
}

func (m *PostMerkle) Put(id int64, data []byte) {
	pos := Post(m.size, m.height, id)
	m.WriteAt(data, pos*hashSize)
}

func (m *PostMerkle) Read(buf []byte, id int64) {
	pos := Post(m.size, m.height, id)
	m.ReadAt(buf, pos*hashSize)
}

func (m *PostMerkle) Build(parallel bool) []byte {
	if parallel {
		return m.parallelbuild()
	} else {
		return m.build()
	}
}

// Iterative post-order depth-first construction of the Merkle tree
// disk access is optimal and forward
func (m *PostMerkle) build() []byte {
	size := m.Size()
	h := sha3.New256()
	hsize := int64(h.Size())

	// pre-allocating the hash stack
	stack := make([][]byte, m.height)
	for i := 0; i < len(stack); i++ {
		stack[i] = make([]byte, hsize)
	}

	var cur int64 = size / 2 // current node (bfs id)
	var count int64 = 0      // post-order id of current node
	var l = 0                // length of the stack

	for count < size {
		if cur >= size/2 { // leaf node
			l++
			h.Reset()
			binary.Write(h, binary.LittleEndian, cur)
			h.Sum(stack[l-1][:0])
			m.WriteAt(stack[l-1], count*hsize)

			for cur&1 == 0 && count < size {
				// we just completed a right node, moving up to the parent
				cur = (cur - 1) >> 1
				count++
				h.Reset()
				h.Write(stack[l-2])
				h.Write(stack[l-1])
				l-- // pop two items, add one item
				h.Sum(stack[l-1][:0])
				m.WriteAt(stack[l-1], count*hsize)
			}
			// we just completed a left node, moving to its sibling
			cur++
			count++
		} else {
			cur = (cur << 1) + 1 // moving to the left child
		}
	}
	return stack[0]
}

func (m *PostMerkle) parallelbuild() (r []byte) {
	var wg sync.WaitGroup
	wg.Add(4)
	go func(id int64) {
		defer wg.Done()
		m.buildsubtree(id)
	}(3)
	go func(id int64) {
		defer wg.Done()
		m.buildsubtree(id)
	}(4)
	go func(id int64) {
		defer wg.Done()
		m.buildsubtree(id)
	}(5)
	go func(id int64) {
		defer wg.Done()
		m.buildsubtree(id)
	}(6)
	wg.Wait()
	buf := make([]byte, hashSize)
	h := sha3.New256()
	m.Read(buf, 3)
	h.Write(buf)
	m.Read(buf, 4)
	h.Write(buf)
	h.Sum(buf[:0])
	m.Put(1, buf)

	h.Reset()
	m.Read(buf, 5)
	h.Write(buf)
	m.Read(buf, 6)
	h.Write(buf)
	h.Sum(buf[:0])
	m.Put(2, buf)

	h.Reset()
	m.Read(buf, 1)
	h.Write(buf)
	m.Read(buf, 2)
	h.Write(buf)
	h.Sum(buf[:0])
	m.Put(0, buf)

	return buf
}

func (m *PostMerkle) buildsubtree(id int64) []byte {
	size := m.Size()
	h := sha3.New256()
	hsize := int64(h.Size())
	height := m.height - Log(id+1)

	// pre-allocating the hash stack
	stack := make([][]byte, height)
	for i := 0; i < len(stack); i++ {
		stack[i] = make([]byte, hsize)
	}

	var cur int64 = id // current node (bfs id)
	for cur < size/2 {
		cur = (cur << 1) + 1
	}
	var count int64 = Post(size, m.height, cur) // post-order id of current node
	var l = 0                                   // length of the stack

	for cur > id {
		if cur >= size/2 { // leaf node
			l++
			h.Reset()
			binary.Write(h, binary.LittleEndian, cur)
			h.Sum(stack[l-1][:0])
			m.WriteAt(stack[l-1], count*hsize)

			for cur&1 == 0 && cur > id {
				// we just completed a right node, moving up to the parent
				cur = (cur - 1) >> 1
				count++
				h.Reset()
				h.Write(stack[l-2])
				h.Write(stack[l-1])
				l-- // pop two items, add one item
				h.Sum(stack[l-1][:0])
				m.WriteAt(stack[l-1], count*hsize)
			}
			if cur <= id {
				break
			}
			// we just completed a left node, moving to its sibling
			cur++
			count++
		} else {
			cur = (cur << 1) + 1 // moving to the left child
		}
	}
	return stack[0]
}

func (m *PostMerkle) Proof(id int64) [][]byte {
	// traversing the tree from the root to the leaf reading the siblings along
	// the path and filling the proof from right to left
	proof := make([][]byte, m.height)
	cur := int64(1)<<uint64(m.height) - 2  // post-order id of the current node along the path, starting from the root
	size := int64(1) << uint64(m.height-1) // size of a subtree of the current node
	mask := size >> 1
	id += 1
	for i := len(proof) - 1; mask > 0; i-- {
		proof[i] = make([]byte, hashSize)
		if mask&id > 0 { // leaf is in the right subtree of current node
			m.ReadAt(proof[i], (cur-size)*hashSize) // reading the left child
			cur -= 1                                // moving to the right subtree
		} else { // leaf is in the left subtree of current node
			m.ReadAt(proof[i], (cur-1)*hashSize) // reading the right child
			cur -= size                          // moving to the left subtree
		}
		size = mask
		mask >>= 1
	}
	proof[0] = make([]byte, hashSize)
	m.ReadAt(proof[0], cur*hashSize)
	return proof
}

func (m *PostMerkle) Proofs(ids []int64) [][][]byte {
	proofs := make([][][]byte, len(ids))
	for i, id := range ids {
		proofs[i] = m.Proof(id)
	}
	return proofs
}