package chain

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

	"basicstack.de/proofedmesh/services/verification-service/internal/canonical"
	"basicstack.de/proofedmesh/services/verification-service/internal/dsse"
	"basicstack.de/proofedmesh/services/verification-service/internal/tlog"
)

// memTree is a test-only Merkle tree mirroring the transparency-log's
// in-memory backend. It lives in the chain tests (not under internal/tlog)
// so the proof verifier under test does not share code with the proof
// producer.
type memTree struct {
	leaves [][]byte
}

func (t *memTree) add(env dsse.Envelope) tlog.Entry {
	leafBytes, _ := canonical.Marshal(env)
	t.leaves = append(t.leaves, leafHash(leafBytes))
	idx := int64(len(t.leaves) - 1)
	hash := "sha256:" + canonical.SHA256Hex(leafBytes)
	return tlog.Entry{
		ContentHash:    hash,
		Envelope:       env,
		LogIndex:       idx,
		InclusionProof: tlog.InclusionProof{},
	}
}

func (t *memTree) refreshProofs(entries ...*tlog.Entry) {
	root := merkleTreeHashT(t.leaves)
	rootStr := "sha256:" + hex.EncodeToString(root)
	for _, e := range entries {
		path := auditPathT(t.leaves, int(e.LogIndex))
		hashes := make([]string, 0, len(path))
		for _, h := range path {
			hashes = append(hashes, "sha256:"+hex.EncodeToString(h))
		}
		e.InclusionProof = tlog.InclusionProof{
			TreeSize:     int64(len(t.leaves)),
			TreeRootHash: rootStr,
			Hashes:       hashes,
		}
	}
}

func leafHash(data []byte) []byte {
	h := sha256.New()
	h.Write([]byte{0x00})
	h.Write(data)
	return h.Sum(nil)
}

func nodeHashT(left, right []byte) []byte {
	h := sha256.New()
	h.Write([]byte{0x01})
	h.Write(left)
	h.Write(right)
	return h.Sum(nil)
}

func merkleTreeHashT(leaves [][]byte) []byte {
	switch len(leaves) {
	case 0:
		s := sha256.Sum256(nil)
		return s[:]
	case 1:
		return append([]byte(nil), leaves[0]...)
	}
	k := largestPow2LessThanT(len(leaves))
	return nodeHashT(merkleTreeHashT(leaves[:k]), merkleTreeHashT(leaves[k:]))
}

func auditPathT(leaves [][]byte, m int) [][]byte {
	n := len(leaves)
	if n <= 1 {
		return nil
	}
	k := largestPow2LessThanT(n)
	if m < k {
		return append(auditPathT(leaves[:k], m), merkleTreeHashT(leaves[k:]))
	}
	return append(auditPathT(leaves[k:], m-k), merkleTreeHashT(leaves[:k]))
}

func largestPow2LessThanT(n int) int {
	k := 1
	for k < n {
		k <<= 1
	}
	return k >> 1
}

func seedBytes(s string) []byte { return []byte(s) }

func sha(b []byte) []byte {
	s := sha256.Sum256(b)
	return s[:]
}