package tlog

import (
	"crypto/ed25519"
	"crypto/rand"
	"encoding/hex"
	"errors"
	"testing"

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

// Local re-implementation of the transparency-log's in-memory Merkle
// backend. Kept here (rather than importing it) so this module stays
// free of cross-module dependencies — the whole point of the proof
// verifier is that it closes against bytes, not against the server
// that produced them.
type tree struct {
	leaves [][]byte
}

func (t *tree) addEnvelope(env dsse.Envelope) (Entry, string) {
	leafBytes, _ := canonical.Marshal(env)
	t.leaves = append(t.leaves, leafHash(leafBytes))
	idx := int64(len(t.leaves) - 1)
	root := merkleTreeHash(t.leaves)
	path := auditPath(t.leaves, int(idx))
	hashes := make([]string, 0, len(path))
	for _, h := range path {
		hashes = append(hashes, "sha256:"+hex.EncodeToString(h))
	}
	hash := "sha256:" + canonical.SHA256Hex(leafBytes)
	return Entry{
		ContentHash: hash,
		Envelope:    env,
		LogIndex:    idx,
		InclusionProof: InclusionProof{
			TreeSize:     int64(len(t.leaves)),
			TreeRootHash: "sha256:" + hex.EncodeToString(root),
			Hashes:       hashes,
		},
	}, hash
}

func merkleTreeHash(leaves [][]byte) []byte {
	switch len(leaves) {
	case 0:
		h := leafHash(nil)
		return h
	case 1:
		return append([]byte(nil), leaves[0]...)
	}
	k := largestPow2LessThan(len(leaves))
	return nodeHash(merkleTreeHash(leaves[:k]), merkleTreeHash(leaves[k:]))
}

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

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

func newEnvelope(t *testing.T, s string) dsse.Envelope {
	t.Helper()
	_, priv, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatalf("gen key: %v", err)
	}
	return dsse.Sign("application/vnd.in-toto+json", []byte(s), "k1", priv)
}

func TestVerifyEntry_SingleLeafTree(t *testing.T) {
	var tr tree
	entry, _ := tr.addEnvelope(newEnvelope(t, `{"i":0}`))
	if err := VerifyEntry(entry); err != nil {
		t.Fatalf("single-leaf: %v", err)
	}
}

func TestVerifyEntry_MultiLeafTree(t *testing.T) {
	var tr tree
	// Append five envelopes; verify every leaf's inclusion proof.
	entries := make([]Entry, 0, 5)
	for i := 0; i < 5; i++ {
		e, _ := tr.addEnvelope(newEnvelope(t, "payload-"+string(rune('a'+i))))
		entries = append(entries, e)
	}
	// Re-fetch the up-to-date proofs for every leaf against the final tree.
	for i := range entries {
		leaf := tr.leaves[i]
		path := auditPath(tr.leaves, i)
		hashes := make([]string, 0, len(path))
		for _, h := range path {
			hashes = append(hashes, "sha256:"+hex.EncodeToString(h))
		}
		entries[i].InclusionProof = InclusionProof{
			TreeSize:     int64(len(tr.leaves)),
			TreeRootHash: "sha256:" + hex.EncodeToString(merkleTreeHash(tr.leaves)),
			Hashes:       hashes,
		}
		_ = leaf
		if err := VerifyEntry(entries[i]); err != nil {
			t.Fatalf("leaf %d: %v", i, err)
		}
	}
}

func TestVerifyEntry_TamperedRoot(t *testing.T) {
	var tr tree
	entry, _ := tr.addEnvelope(newEnvelope(t, `{"i":0}`))
	tr.addEnvelope(newEnvelope(t, `{"i":1}`))
	// Re-fetch leaf 0's proof against the 2-leaf tree.
	path := auditPath(tr.leaves, 0)
	hashes := make([]string, 0, len(path))
	for _, h := range path {
		hashes = append(hashes, "sha256:"+hex.EncodeToString(h))
	}
	entry.InclusionProof = InclusionProof{
		TreeSize:     2,
		TreeRootHash: "sha256:" + hex.EncodeToString(merkleTreeHash(tr.leaves)),
		Hashes:       hashes,
	}
	// Flip a bit in the stated root — proof must fail.
	rootBytes, _ := hex.DecodeString(entry.InclusionProof.TreeRootHash[len("sha256:"):])
	rootBytes[0] ^= 0xff
	entry.InclusionProof.TreeRootHash = "sha256:" + hex.EncodeToString(rootBytes)
	err := VerifyEntry(entry)
	if !errors.Is(err, ErrInclusionProofInvalid) {
		t.Fatalf("want ErrInclusionProofInvalid, got %v", err)
	}
}

func TestVerifyEntry_ContentHashMismatch(t *testing.T) {
	var tr tree
	entry, _ := tr.addEnvelope(newEnvelope(t, `{"i":0}`))
	// Flip a bit in the envelope payload after the proof was captured.
	entry.Envelope.Payload = "aaaa"
	err := VerifyEntry(entry)
	if !errors.Is(err, ErrInclusionProofMalformed) {
		t.Fatalf("want ErrInclusionProofMalformed, got %v", err)
	}
}

func TestVerifyEntry_IndexOutOfRange(t *testing.T) {
	var tr tree
	entry, _ := tr.addEnvelope(newEnvelope(t, `{"i":0}`))
	entry.LogIndex = 7
	entry.InclusionProof.TreeSize = 7
	err := VerifyEntry(entry)
	if !errors.Is(err, ErrInclusionProofMalformed) {
		t.Fatalf("want ErrInclusionProofMalformed, got %v", err)
	}
}