package backend

import (
	"bytes"
	"context"
	"encoding/hex"
	"errors"
	"strings"
	"testing"
	"time"
)

func sampleEnvelope(payload, sig string) Envelope {
	return Envelope{
		PayloadType: "application/vnd.in-toto+json",
		Payload:     payload,
		Signatures:  []EnvelopeSignature{{KeyID: "k1", Sig: sig}},
	}
}

func mustHash(t *testing.T, env Envelope) string {
	t.Helper()
	h, err := EnvelopeContentHash(env)
	if err != nil {
		t.Fatalf("hash: %v", err)
	}
	return h
}

func TestMemory_SubmitAndIdempotency(t *testing.T) {
	m := NewMemory().WithClock(func() time.Time {
		return time.Date(2026, 10, 25, 12, 0, 0, 0, time.UTC)
	})
	env := sampleEnvelope("aGVsbG8", "c2ln")
	ch := mustHash(t, env)

	r1, fresh, err := m.Submit(context.Background(), env, ch)
	if err != nil {
		t.Fatalf("first submit: %v", err)
	}
	if !fresh {
		t.Errorf("first submit should be fresh")
	}
	if r1.LogIndex != 0 {
		t.Errorf("first log_index: want 0, got %d", r1.LogIndex)
	}
	if r1.ContentHash != ch {
		t.Errorf("content_hash echoed wrong: %q vs %q", r1.ContentHash, ch)
	}
	if r1.InclusionProof.TreeSize != 1 {
		t.Errorf("tree_size: %d", r1.InclusionProof.TreeSize)
	}
	if len(r1.InclusionProof.Hashes) != 0 {
		t.Errorf("single-leaf audit path should be empty, got %d entries", len(r1.InclusionProof.Hashes))
	}

	r2, fresh2, err := m.Submit(context.Background(), env, ch)
	if err != nil {
		t.Fatalf("second submit: %v", err)
	}
	if fresh2 {
		t.Errorf("second submit should NOT be fresh (idempotent repeat)")
	}
	if r2.LogIndex != r1.LogIndex {
		t.Errorf("idempotent log_index drifted: %d -> %d", r1.LogIndex, r2.LogIndex)
	}
	if !r2.LoggedAt.Equal(r1.LoggedAt) {
		t.Errorf("logged_at should be stable across idempotent repeats")
	}
	if r2.InclusionProof.TreeRootHash != r1.InclusionProof.TreeRootHash {
		t.Errorf("tree_root drifted without new entries: %q vs %q", r1.InclusionProof.TreeRootHash, r2.InclusionProof.TreeRootHash)
	}
}

func TestMemory_ContentHashMismatch(t *testing.T) {
	m := NewMemory()
	env := sampleEnvelope("aGVsbG8", "c2ln")
	wrong := "sha256:" + strings.Repeat("0", 64)
	_, _, err := m.Submit(context.Background(), env, wrong)
	if err == nil {
		t.Fatalf("expected ErrContentHashMismatch, got nil")
	}
	if err != ErrContentHashMismatch {
		t.Fatalf("wrong error: %v", err)
	}
}

func TestMemory_MultipleEntriesGrowRoot(t *testing.T) {
	m := NewMemory()
	ctx := context.Background()
	var lastRoot string
	for i, payload := range []string{"a", "b", "c", "d", "e"} {
		env := sampleEnvelope(payload, "s")
		ch := mustHash(t, env)
		r, fresh, err := m.Submit(ctx, env, ch)
		if err != nil {
			t.Fatalf("submit %d: %v", i, err)
		}
		if !fresh {
			t.Errorf("submit %d should be fresh", i)
		}
		if r.LogIndex != int64(i) {
			t.Errorf("submit %d log_index: want %d got %d", i, i, r.LogIndex)
		}
		if r.InclusionProof.TreeSize != int64(i+1) {
			t.Errorf("submit %d tree_size: want %d got %d", i, i+1, r.InclusionProof.TreeSize)
		}
		if r.InclusionProof.TreeRootHash == lastRoot {
			t.Errorf("tree_root did not change after new entry %d", i)
		}
		lastRoot = r.InclusionProof.TreeRootHash
	}
}

// TestMemory_InclusionProofVerifiesAgainstRoot cross-checks that for a tree
// of several entries (including a non-power-of-two size), the audit path
// returned alongside every leaf verifies against the advertised tree root
// per RFC 6962 §2.1.1. If this breaks, our Merkle code is wrong.
func TestMemory_InclusionProofVerifiesAgainstRoot(t *testing.T) {
	for _, n := range []int{1, 2, 3, 4, 5, 7, 8, 13} {
		m := NewMemory()
		ctx := context.Background()
		envs := make([]Envelope, n)
		hashes := make([]string, n)
		for i := range envs {
			envs[i] = sampleEnvelope(string(rune('A'+i)), "s")
			hashes[i] = mustHash(t, envs[i])
			if _, _, err := m.Submit(ctx, envs[i], hashes[i]); err != nil {
				t.Fatalf("n=%d seed %d: %v", n, i, err)
			}
		}
		for i := range envs {
			r, _, err := m.Submit(ctx, envs[i], hashes[i])
			if err != nil {
				t.Fatalf("n=%d fetch leaf %d: %v", n, i, err)
			}
			canonical, err := canonicalJSON(envs[i])
			if err != nil {
				t.Fatalf("canonical: %v", err)
			}
			leaf := leafHash(canonical)
			rootBytes, err := decodePrefixedHex(r.InclusionProof.TreeRootHash)
			if err != nil {
				t.Fatalf("decode root: %v", err)
			}
			var proof [][]byte
			for _, h := range r.InclusionProof.Hashes {
				b, err := decodePrefixedHex(h)
				if err != nil {
					t.Fatalf("decode sibling: %v", err)
				}
				proof = append(proof, b)
			}
			if err := verifyInclusionRFC6962(int64(i), int64(n), leaf, proof, rootBytes); err != nil {
				t.Errorf("n=%d leaf %d: audit path does not verify: %v", n, i, err)
			}
		}
	}
}

// verifyInclusionRFC6962 implements the RFC 6962 §2.1.1 audit path
// verification algorithm. Returns nil if proof proves leafHash is at
// leafIndex in a tree of size treeSize with the given root.
func verifyInclusionRFC6962(leafIndex, treeSize int64, leaf []byte, proof [][]byte, root []byte) error {
	if leafIndex >= treeSize {
		return errors.New("leaf index out of range")
	}
	fn, sn := leafIndex, treeSize-1
	r := append([]byte(nil), leaf...)
	for _, p := range proof {
		if sn == 0 {
			return errors.New("extra proof element")
		}
		if fn&1 == 1 || fn == sn {
			r = nodeHash(p, r)
			if fn&1 == 0 {
				for fn != 0 && fn&1 == 0 {
					fn >>= 1
					sn >>= 1
				}
			}
		} else {
			r = nodeHash(r, p)
		}
		fn >>= 1
		sn >>= 1
	}
	if sn != 0 {
		return errors.New("incomplete proof")
	}
	if !bytes.Equal(r, root) {
		return errors.New("root mismatch")
	}
	return nil
}

func decodePrefixedHex(s string) ([]byte, error) {
	const prefix = "sha256:"
	if len(s) <= len(prefix) || s[:len(prefix)] != prefix {
		return nil, errors.New("hash missing sha256: prefix")
	}
	return hex.DecodeString(s[len(prefix):])
}