package checkpoint

import (
	"context"
	"crypto/ed25519"
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"github.com/Fuwn/plutia/internal/storage"
	"github.com/Fuwn/plutia/internal/types"
	"os"
	"path/filepath"
	"testing"
)

func TestBuildAndStoreCheckpoint(t *testing.T) {
	tmp := t.TempDir()
	store, err := storage.OpenPebble(tmp)

	if err != nil {
		t.Fatalf("open pebble: %v", err)
	}

	defer store.Close()

	pub, priv, err := ed25519.GenerateKey(rand.Reader)

	if err != nil {
		t.Fatalf("generate key: %v", err)
	}

	keyPath := filepath.Join(tmp, "mirror.key")

	if err := os.WriteFile(keyPath, []byte(base64.RawURLEncoding.EncodeToString(priv)), 0o600); err != nil {
		t.Fatalf("write key: %v", err)
	}

	mgr := NewManager(store, tmp, keyPath)
	states := []types.StateV1{
		{Version: 1, DID: "did:plc:a", ChainTipHash: "tip-a"},
		{Version: 1, DID: "did:plc:b", ChainTipHash: "tip-b"},
	}
	cp, err := mgr.BuildAndStore(42, states, []string{"abc", "def"})

	if err != nil {
		t.Fatalf("build checkpoint: %v", err)
	}

	if cp.Signature == "" || cp.CheckpointHash == "" {
		t.Fatalf("expected signed checkpoint")
	}

	payload, err := marshalCheckpointPayload(cp)

	if err != nil {
		t.Fatalf("marshal payload: %v", err)
	}

	rawSig, err := base64.RawURLEncoding.DecodeString(cp.Signature)

	if err != nil {
		t.Fatalf("decode signature: %v", err)
	}

	if !ed25519.Verify(pub, payload, rawSig) {
		t.Fatalf("signature verification failed")
	}

	stored, ok, err := store.GetCheckpoint(cp.Sequence)

	if err != nil || !ok {
		t.Fatalf("stored checkpoint missing: ok=%v err=%v", ok, err)
	}

	if stored.CheckpointHash != cp.CheckpointHash {
		t.Fatalf("checkpoint hash mismatch")
	}
}

func TestCheckpointRootStability(t *testing.T) {
	tmp := t.TempDir()
	store, err := storage.OpenPebble(tmp)

	if err != nil {
		t.Fatalf("open pebble: %v", err)
	}

	defer store.Close()

	if err := store.PutState(types.StateV1{Version: 1, DID: "did:plc:b", ChainTipHash: "tip-b"}); err != nil {
		t.Fatalf("put state b: %v", err)
	}

	if err := store.PutState(types.StateV1{Version: 1, DID: "did:plc:a", ChainTipHash: "tip-a"}); err != nil {
		t.Fatalf("put state a: %v", err)
	}

	_, priv, err := ed25519.GenerateKey(rand.Reader)

	if err != nil {
		t.Fatalf("generate key: %v", err)
	}

	keyPath := filepath.Join(tmp, "mirror.key")

	if err := os.WriteFile(keyPath, []byte(base64.RawURLEncoding.EncodeToString(priv)), 0o600); err != nil {
		t.Fatalf("write key: %v", err)
	}

	mgr := NewManager(store, tmp, keyPath)
	cp1, err := mgr.BuildAndStoreFromStore(100, []string{"abc"})

	if err != nil {
		t.Fatalf("build checkpoint 1: %v", err)
	}

	cp2, err := mgr.BuildAndStoreFromStore(200, []string{"abc"})

	if err != nil {
		t.Fatalf("build checkpoint 2: %v", err)
	}

	if cp1.DIDMerkleRoot != cp2.DIDMerkleRoot {
		t.Fatalf("did root changed for identical state: %s vs %s", cp1.DIDMerkleRoot, cp2.DIDMerkleRoot)
	}

	if cp1.BlockMerkleRoot != cp2.BlockMerkleRoot {
		t.Fatalf("block root changed for identical block set: %s vs %s", cp1.BlockMerkleRoot, cp2.BlockMerkleRoot)
	}
}

func TestBuildDIDProofAtCheckpointHonorsContextCancellation(t *testing.T) {
	tmp := t.TempDir()
	store, err := storage.OpenPebble(tmp)

	if err != nil {
		t.Fatalf("open pebble: %v", err)
	}

	defer store.Close()

	_, priv, err := ed25519.GenerateKey(rand.Reader)

	if err != nil {
		t.Fatalf("generate key: %v", err)
	}

	keyPath := filepath.Join(tmp, "mirror.key")

	if err := os.WriteFile(keyPath, []byte(base64.RawURLEncoding.EncodeToString(priv)), 0o600); err != nil {
		t.Fatalf("write key: %v", err)
	}

	checkpointsDir := filepath.Join(tmp, "checkpoints")

	if err := os.MkdirAll(checkpointsDir, 0o755); err != nil {
		t.Fatalf("mkdir checkpoints: %v", err)
	}

	leaves := make([]types.DIDLeaf, 500)

	for i := range leaves {
		leaves[i] = types.DIDLeaf{
			DID:          "did:plc:test-" + string(rune('a'+(i%26))),
			ChainTipHash: "tip",
		}
	}

	snapshot := types.CheckpointStateSnapshotV1{
		Version:   1,
		Sequence:  10,
		CreatedAt: "2026-02-26T00:00:00Z",
		Leaves:    leaves,
	}
	b, err := json.Marshal(snapshot)

	if err != nil {
		t.Fatalf("marshal snapshot: %v", err)
	}

	if err := os.WriteFile(filepath.Join(checkpointsDir, "00000000000000000010.state.json"), b, 0o644); err != nil {
		t.Fatalf("write snapshot: %v", err)
	}

	mgr := NewManager(store, tmp, keyPath)
	ctx, cancel := context.WithCancel(context.Background())

	cancel()

	if _, _, _, err := mgr.BuildDIDProofAtCheckpoint(ctx, "did:plc:test-a", "tip", 10); err == nil {
		t.Fatalf("expected context cancellation error")
	}
}