package store

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"testing"
)

func TestMemoryModuleSnapshotAtomicConcurrentRMW(t *testing.T) {
	m := NewMemory()
	m.SeedDemo()
	tenant, _, _, _, _ := m.DemoIDs()
	mod, err := m.CreateModule(tenant, &Module{Type: "CDN_CIDRS", Name: "atomic", Enabled: true})
	if err != nil {
		t.Fatal(err)
	}
	const writers = 12
	var wg sync.WaitGroup
	errCh := make(chan error, writers)
	for i := 0; i < writers; i++ {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			errCh <- m.UpdateModulePrefixSnapshot(context.Background(), tenant, mod.ID, func(current *ModulePrefixSnapshot) (*ModulePrefixSnapshot, error) {
				var rows []PrefixRow
				if current != nil {
					rows = append(rows, current.Prefixes...)
				}
				rows = append(rows, PrefixRow{Prefix: fmt.Sprintf("192.0.2.%d/32", i)})
				return &ModulePrefixSnapshot{InputHash: "atomic", Prefixes: rows}, nil
			})
		}(i)
	}
	wg.Wait()
	close(errCh)
	for err := range errCh {
		if err != nil {
			t.Fatal(err)
		}
	}
	snap, ok, err := m.GetModulePrefixSnapshot(context.Background(), tenant, mod.ID)
	if err != nil || !ok || len(snap.Prefixes) != writers {
		t.Fatalf("lost concurrent update: %+v ok=%v err=%v", snap, ok, err)
	}
	if got, err := m.GetModule(tenant, mod.ID); err != nil || got == nil || got.InputHash != snap.InputHash {
		t.Fatalf("module input hash differs from snapshot: %+v %+v", got, snap)
	}
}

func TestMemoryModuleSnapshotAtomicCallbackErrorRollsBack(t *testing.T) {
	m := NewMemory()
	m.SeedDemo()
	tenant, _, _, _, _ := m.DemoIDs()
	mod, err := m.CreateModule(tenant, &Module{Type: "CDN_CIDRS", Name: "rollback", Enabled: true})
	if err != nil {
		t.Fatal(err)
	}
	ctx := context.Background()
	if err := m.SetModulePrefixSnapshot(ctx, tenant, mod.ID, "before", []PrefixRow{{Prefix: "192.0.2.0/24"}}); err != nil {
		t.Fatal(err)
	}
	want := errors.New("update failed")
	err = m.UpdateModulePrefixSnapshot(ctx, tenant, mod.ID, func(current *ModulePrefixSnapshot) (*ModulePrefixSnapshot, error) {
		current.Prefixes[0].Prefix = "198.51.100.0/24"
		return nil, want
	})
	if !errors.Is(err, want) {
		t.Fatalf("callback error: %v", err)
	}
	snap, ok, err := m.GetModulePrefixSnapshot(ctx, tenant, mod.ID)
	if err != nil || !ok || snap.InputHash != "before" || snap.Prefixes[0].Prefix != "192.0.2.0/24" {
		t.Fatalf("partial snapshot visible: %+v ok=%v err=%v", snap, ok, err)
	}
}