// KEKS/CM -- KEKS-encoded cryptographic messages
// Copyright (C) 2024-2025 Sergey Matveev <stargrave@stargrave.org>
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Lesser General Public License as
// published by the Free Software Foundation, version 3 of the License.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this program.  If not, see <http://www.gnu.org/licenses/>.

package merkle

import (
	"errors"
	"hash"
	"io"
	"os"

	"golang.org/x/sys/unix"
)

const (
	maxDepth        = 64
	Leaf            = "LEAF"
	Node            = "NODE"
	DefaultChunkLen = 128 * 1024
)

type job struct {
	bufReady chan struct{}
	hshReady chan struct{}
	buf      []byte
	hsh      []byte
}

// Merkle-tree based parallelised hasher. You must use either io.Writer
// interface to feed the data, or call DoReadFrom() method only once, or
// call Mmap() method.
type Hasher struct {
	ready     chan *job
	dones     chan *job
	finished  chan struct{}
	leafNew   func() hash.Hash
	leafReset func(hash.Hash) hash.Hash
	nodeReset func(hash.Hash) hash.Hash
	nodeHash  hash.Hash
	pr        *io.PipeReader
	pw        *io.PipeWriter
	hashes    [2 * maxDepth][]byte
	frees     [2 * maxDepth]bool
	count     int
	workers   int
	chunkLen  int
	wasCalled bool
}

func (h *Hasher) Size() int {
	return h.nodeHash.Size()
}

func (h *Hasher) BlockSize() int {
	return h.nodeHash.BlockSize()
}

// Create new Merkle-tree based hasher. leafNew/nodeNew specify
// functions that will create leaf/node hasher correspondingly.
// leafReset/nodeReset are functions to reset them, to avoid
// possibly relatively expensive hash creation call.
func NewHasher(
	leafNew, nodeNew func() hash.Hash,
	leafReset, nodeReset func(hash.Hash) hash.Hash,
	chunkLen, workers int,
) *Hasher {
	h := Hasher{
		leafNew:   leafNew,
		nodeHash:  nodeNew(),
		leafReset: leafReset,
		nodeReset: nodeReset,
		workers:   workers,
		chunkLen:  chunkLen,
	}
	hashSize := h.Size()
	for i := range 2 * maxDepth {
		h.hashes[i] = make([]byte, hashSize)
		h.frees[i] = true
	}
	h.prepare()
	return &h
}

func (h *Hasher) prepare() {
	h.finished = make(chan struct{})
	h.ready = make(chan *job, h.workers)
	h.dones = make(chan *job, h.workers)
	for range h.workers {
		go h.worker()
	}
	go h.aggregator()
}

func (h *Hasher) Reset() {
	h.pw.Close()
	<-h.finished
	for i := range 2 * maxDepth {
		h.frees[i] = true
	}
	h.prepare()
}

func (h *Hasher) get(l int) []byte {
	if l >= maxDepth {
		panic("too deep")
	}
	i := l * 2
	if !h.frees[i] {
		i++
	}
	h.frees[i] = false
	return h.hashes[i]
}

func (h *Hasher) fold() {
	var err error
	for l := range maxDepth {
		if h.frees[l*2+0] || h.frees[l*2+1] {
			continue
		}
		h.nodeHash = h.nodeReset(h.nodeHash)
		if _, err = h.nodeHash.Write(h.hashes[l*2+0]); err != nil {
			panic(err)
		}
		if _, err = h.nodeHash.Write(h.hashes[l*2+1]); err != nil {
			panic(err)
		}
		h.nodeHash.Sum(h.get(l + 1)[:0])
		h.frees[l*2+0] = true
		h.frees[l*2+1] = true
		h.count--
	}
}

func (h *Hasher) makePipe() {
	h.pr, h.pw = io.Pipe()
	go func() {
		h.DoReadFrom(h.pr)
		h.pr.Close()
	}()
}

func (h *Hasher) Write(p []byte) (int, error) {
	if h.pr == nil {
		h.makePipe()
	}
	return h.pw.Write(p)
}

func (h *Hasher) DoReadFrom(r io.Reader) (total int64, err error) {
	if h.wasCalled {
		panic("must be called only once")
	}
	h.wasCalled = true
	defer close(h.dones)
	var j *job
	var n int
	var eof bool
	for !eof {
		j = <-h.ready
		n, err = io.ReadFull(r, j.buf)
		total += int64(n)
		if err != nil {
			if err != io.ErrUnexpectedEOF {
				return
			}
			err = nil
			eof = true
			j.buf = j.buf[:n]
		}
		if n == 0 {
			continue
		}
		j.bufReady <- struct{}{}
		h.dones <- j
	}
	return
}

func (h *Hasher) worker() {
	j := job{
		buf:      make([]byte, h.chunkLen),
		hsh:      make([]byte, h.Size()),
		bufReady: make(chan struct{}),
		hshReady: make(chan struct{}),
	}
	h.ready <- &j
	hasher := h.leafNew()
	var err error
	for {
		<-j.bufReady
		hasher = h.leafReset(hasher)
		if _, err = hasher.Write(j.buf); err != nil {
			panic(err)
		}
		hasher.Sum(j.hsh[:0])
		j.hshReady <- struct{}{}
	}
}

func (h *Hasher) aggregator() {
	for j := range h.dones {
		<-j.hshReady
		copy(h.get(0), j.hsh)
		h.ready <- j
		h.count++
		h.fold()
	}
	close(h.finished)
}

func (h *Hasher) Sum(b []byte) []byte {
	if !h.wasCalled {
		return append(b, h.leafNew().Sum(nil)...)
	}
	if h.pw != nil {
		h.pw.Close()
	}
	<-h.finished
	for l := 0; l < maxDepth; l++ {
		if h.count == 1 {
			for ; l < maxDepth; l++ {
				if !h.frees[l*2+0] {
					return append(b, h.hashes[l*2+0]...)
				}
			}
		}
		if !h.frees[l*2+0] {
			copy(h.get(l+1), h.hashes[l*2+0])
			h.frees[l*2+0] = true
			h.fold()
		}
	}
	panic("did not reach the end")
}

func (h *Hasher) Mmap(fn string) error {
	fd, err := os.Open(fn)
	if err != nil {
		return err
	}
	defer fd.Close()
	var size int64
	{
		var fi os.FileInfo
		fi, err = fd.Stat()
		if err != nil {
			return err
		}
		size = fi.Size()
	}
	if size == 0 {
		return nil
	}
	if size < 0 {
		return errors.New("negative size")
	}
	if size != int64(int(size)) {
		return errors.New("file is too large")
	}
	data, err := unix.Mmap(int(fd.Fd()), 0, int(size), unix.PROT_READ, unix.MAP_SHARED)
	if err != nil {
		return err
	}

	h.wasCalled = true
	var j *job
	for i := 0; i < len(data); i += h.chunkLen {
		j = <-h.ready
		j.buf = data[i:min(i+h.chunkLen, len(data))]
		if len(j.buf) == 0 {
			break
		}
		j.bufReady <- struct{}{}
		h.dones <- j
	}
	close(h.dones)
	<-h.finished
	unix.Munmap(data)
	return nil
}
