package workerruntime

import (
	"context"
	"crypto/sha256"
	"errors"
	"testing"
	"time"
)

func TestMemoryStoreRequiresAcknowledgedSnapshotForRuntime(t *testing.T) {
	now := time.Date(2026, 7, 31, 9, 0, 0, 0, time.UTC)
	store := newRuntimeStore(t, &now)
	ctx := context.Background()
	openControlSession(t, store, "session-a", time.Minute)
	epoch, err := store.CurrentOwnershipEpoch(ctx)
	if err != nil {
		t.Fatalf("CurrentOwnershipEpoch(): %v", err)
	}
	report := controlReport(now, "session-a", 1, 7, epoch)
	if err := store.ReplaceRuntime(ctx, report, time.Minute); !errors.Is(err, ErrSnapshotMismatch) {
		t.Fatalf("ReplaceRuntime(before ACK) error = %v, want ErrSnapshotMismatch", err)
	}
	reference := controlReference(7, epoch, "snapshot-7")
	if err := store.RecordIssuedSnapshot(ctx, reference, time.Minute); err != nil {
		t.Fatalf("RecordIssuedSnapshot(): %v", err)
	}
	if err := store.AcknowledgeSnapshot(ctx, SnapshotAcknowledgement{
		WorkerID: "worker-a", SessionID: "session-a", Reference: reference, Applied: true,
	}, time.Minute); err != nil {
		t.Fatalf("AcknowledgeSnapshot(): %v", err)
	}
	if err := store.ReplaceRuntime(ctx, report, time.Minute); err != nil {
		t.Fatalf("ReplaceRuntime(after ACK): %v", err)
	}
}

func TestMemoryStoreNegativeAcknowledgementFencesDelayedRuntime(t *testing.T) {
	now := time.Date(2026, 7, 31, 9, 0, 0, 0, time.UTC)
	store := newRuntimeStore(t, &now)
	ctx := context.Background()
	openControlSession(t, store, "session-a", time.Minute)
	epoch, err := store.CurrentOwnershipEpoch(ctx)
	if err != nil {
		t.Fatalf("CurrentOwnershipEpoch(): %v", err)
	}
	first := controlReference(7, epoch, "snapshot-7")
	if err := store.RecordIssuedSnapshot(ctx, first, time.Minute); err != nil {
		t.Fatalf("RecordIssuedSnapshot(first): %v", err)
	}
	if err := store.AcknowledgeSnapshot(ctx, SnapshotAcknowledgement{WorkerID: "worker-a", SessionID: "session-a", Reference: first, Applied: true}, time.Minute); err != nil {
		t.Fatalf("AcknowledgeSnapshot(first): %v", err)
	}
	if err := store.ReplaceRuntime(ctx, controlReport(now, "session-a", 1, 7, epoch), time.Minute); err != nil {
		t.Fatalf("ReplaceRuntime(first): %v", err)
	}

	second := controlReference(8, epoch, "snapshot-8")
	if err := store.RecordIssuedSnapshot(ctx, second, time.Minute); err != nil {
		t.Fatalf("RecordIssuedSnapshot(second): %v", err)
	}
	if err := store.AcknowledgeSnapshot(ctx, SnapshotAcknowledgement{WorkerID: "worker-a", SessionID: "session-a", Reference: second, Applied: false, ErrorCode: "apply_failed"}, time.Minute); err != nil {
		t.Fatalf("AcknowledgeSnapshot(negative): %v", err)
	}
	if err := store.ReplaceRuntime(ctx, controlReport(now, "session-a", 2, 7, epoch), time.Minute); !errors.Is(err, ErrSnapshotMismatch) {
		t.Fatalf("ReplaceRuntime(delayed): %v, want ErrSnapshotMismatch", err)
	}
	if err := store.AcknowledgeSnapshot(ctx, SnapshotAcknowledgement{WorkerID: "worker-a", SessionID: "session-a", Reference: second, Applied: true}, time.Minute); err != nil {
		t.Fatalf("AcknowledgeSnapshot(recover): %v", err)
	}
	if err := store.ReplaceRuntime(ctx, controlReport(now, "session-a", 2, 8, epoch), time.Minute); err != nil {
		t.Fatalf("ReplaceRuntime(recovered): %v", err)
	}
}

func TestMemoryStoreAcknowledgementReplayPreservesRuntimeFence(t *testing.T) {
	now := time.Date(2026, 7, 31, 9, 0, 0, 0, time.UTC)
	store := newRuntimeStore(t, &now)
	ctx := context.Background()
	openControlSession(t, store, "session-a", time.Minute)
	epoch, err := store.CurrentOwnershipEpoch(ctx)
	if err != nil {
		t.Fatalf("CurrentOwnershipEpoch(): %v", err)
	}
	reference := controlReference(7, epoch, "snapshot-7")
	if err := store.RecordIssuedSnapshot(ctx, reference, time.Minute); err != nil {
		t.Fatalf("RecordIssuedSnapshot(): %v", err)
	}
	ack := SnapshotAcknowledgement{WorkerID: "worker-a", SessionID: "session-a", Reference: reference, Applied: true}
	if err := store.AcknowledgeSnapshot(ctx, ack, time.Minute); err != nil {
		t.Fatalf("AcknowledgeSnapshot(first): %v", err)
	}
	if err := store.ReplaceRuntime(ctx, controlReport(now, "session-a", 100, 7, epoch), time.Minute); err != nil {
		t.Fatalf("ReplaceRuntime(): %v", err)
	}
	if err := store.AcknowledgeSnapshot(ctx, ack, time.Minute); err != nil {
		t.Fatalf("AcknowledgeSnapshot(replay): %v", err)
	}
	got, err := store.ReadRuntime(ctx, []OwnedProxy{{ProxyID: "proxy-a", WorkerID: "worker-a", OwnershipEpoch: epoch}})
	if err != nil || len(got) != 1 || !got[0].Fresh {
		t.Fatalf("ReadRuntime(after ACK replay) = %+v, %v", got, err)
	}
	if err := store.ReplaceRuntime(ctx, controlReport(now, "session-a", 99, 7, epoch), time.Minute); !errors.Is(err, ErrStaleReport) {
		t.Fatalf("ReplaceRuntime(stale): %v, want ErrStaleReport", err)
	}
}

func TestMemoryStoreRejectsConflictingSnapshotReference(t *testing.T) {
	now := time.Date(2026, 7, 31, 9, 0, 0, 0, time.UTC)
	store := newRuntimeStore(t, &now)
	ctx := context.Background()
	epoch, err := store.CurrentOwnershipEpoch(ctx)
	if err != nil {
		t.Fatalf("CurrentOwnershipEpoch(): %v", err)
	}
	if err := store.RecordIssuedSnapshot(ctx, controlReference(7, epoch, "first"), time.Minute); err != nil {
		t.Fatalf("RecordIssuedSnapshot(first): %v", err)
	}
	if err := store.RecordIssuedSnapshot(ctx, controlReference(7, epoch, "second"), time.Minute); !errors.Is(err, ErrConflictingSnapshotReference) {
		t.Fatalf("RecordIssuedSnapshot(conflict): %v, want ErrConflictingSnapshotReference", err)
	}
}

func TestMemoryStoreReplacesSparseRuntimeAndClearsMissingCounters(t *testing.T) {
	now := time.Date(2026, 7, 30, 15, 0, 0, 0, time.UTC)
	store := newRuntimeStore(t, &now)
	ctx := context.Background()
	registerRuntimeSession(t, store, "session-a", time.Minute)
	report := Report{
		WorkerID: "worker-a", SessionID: "session-a", Sequence: 1,
		SnapshotVersion: 3, OwnershipEpoch: 9, ObservedAt: now,
		Counters: []Counter{{ProxyID: "proxy-a", Active: 2, Reserved: 1}},
	}
	if err := store.ReplaceRuntime(ctx, report, time.Minute); err != nil {
		t.Fatalf("ReplaceRuntime(first): %v", err)
	}
	got, err := store.ReadRuntime(ctx, []OwnedProxy{{ProxyID: "proxy-a", WorkerID: "worker-a", OwnershipEpoch: 9}})
	if err != nil || len(got) != 1 || got[0] != (Snapshot{ProxyID: "proxy-a", Active: 2, Reserved: 1, Fresh: true}) {
		t.Fatalf("ReadRuntime(first) = %+v, %v", got, err)
	}

	report.Sequence = 2
	report.ObservedAt = now.Add(time.Second)
	report.Counters = nil
	if err := store.ReplaceRuntime(ctx, report, time.Minute); err != nil {
		t.Fatalf("ReplaceRuntime(empty): %v", err)
	}
	got, err = store.ReadRuntime(ctx, []OwnedProxy{{ProxyID: "proxy-a", WorkerID: "worker-a", OwnershipEpoch: 9}})
	if err != nil || len(got) != 1 || got[0] != (Snapshot{ProxyID: "proxy-a", Fresh: true}) {
		t.Fatalf("ReadRuntime(empty) = %+v, %v", got, err)
	}
}

func TestMemoryStoreFencesSessionsAndReportSequence(t *testing.T) {
	now := time.Date(2026, 7, 30, 15, 0, 0, 0, time.UTC)
	store := newRuntimeStore(t, &now)
	ctx := context.Background()
	registerRuntimeSession(t, store, "session-a", time.Minute)
	report := Report{
		WorkerID: "worker-a", SessionID: "session-a", Sequence: 2,
		SnapshotVersion: 3, OwnershipEpoch: 9, ObservedAt: now,
		Counters: []Counter{{ProxyID: "proxy-a", Active: 1}},
	}
	if err := store.ReplaceRuntime(ctx, report, time.Minute); err != nil {
		t.Fatalf("ReplaceRuntime(first): %v", err)
	}
	if err := store.ReplaceRuntime(ctx, report, time.Minute); err != nil {
		t.Fatalf("ReplaceRuntime(idempotent): %v", err)
	}
	conflict := report
	conflict.Counters = []Counter{{ProxyID: "proxy-a", Active: 2}}
	if err := store.ReplaceRuntime(ctx, conflict, time.Minute); !errors.Is(err, ErrConflictingReport) {
		t.Fatalf("ReplaceRuntime(conflict) error = %v", err)
	}
	stale := report
	stale.Sequence = 1
	if err := store.ReplaceRuntime(ctx, stale, time.Minute); !errors.Is(err, ErrStaleReport) {
		t.Fatalf("ReplaceRuntime(stale) error = %v", err)
	}
	registerRuntimeSession(t, store, "session-b", time.Minute)
	if err := store.ReplaceRuntime(ctx, report, time.Minute); !errors.Is(err, ErrStaleSession) {
		t.Fatalf("ReplaceRuntime(old session) error = %v", err)
	}
}

func TestMemoryStoreFailsClosedForExpiredOrOlderOwnershipReport(t *testing.T) {
	now := time.Date(2026, 7, 30, 15, 0, 0, 0, time.UTC)
	store := newRuntimeStore(t, &now)
	ctx := context.Background()
	registerRuntimeSession(t, store, "session-a", time.Minute)
	if err := store.ReplaceRuntime(ctx, Report{
		WorkerID: "worker-a", SessionID: "session-a", Sequence: 1,
		SnapshotVersion: 3, OwnershipEpoch: 9, ObservedAt: now,
	}, time.Minute); err != nil {
		t.Fatalf("ReplaceRuntime(): %v", err)
	}
	queries := []OwnedProxy{
		{ProxyID: "proxy-a", WorkerID: "worker-a", OwnershipEpoch: 10},
		{ProxyID: "proxy-b", WorkerID: "worker-a", OwnershipEpoch: 9},
	}
	got, err := store.ReadRuntime(ctx, queries)
	if err != nil || got[0].Fresh || !got[1].Fresh {
		t.Fatalf("ReadRuntime(ownership fence) = %+v, %v", got, err)
	}
	now = now.Add(time.Minute)
	got, err = store.ReadRuntime(ctx, queries[1:])
	if err != nil || len(got) != 1 || got[0].Fresh {
		t.Fatalf("ReadRuntime(expired) = %+v, %v", got, err)
	}
}

func TestMemoryStoreRejectsRuntimeBeyondAcknowledgedSnapshot(t *testing.T) {
	now := time.Date(2026, 7, 30, 15, 0, 0, 0, time.UTC)
	store := newRuntimeStore(t, &now)
	registerRuntimeSession(t, store, "session-a", time.Minute)
	report := Report{
		WorkerID: "worker-a", SessionID: "session-a", Sequence: 1,
		SnapshotVersion: 4, OwnershipEpoch: 9, ObservedAt: now,
	}
	if err := store.ReplaceRuntime(context.Background(), report, time.Minute); !errors.Is(err, ErrSnapshotMismatch) {
		t.Fatalf("ReplaceRuntime(ahead snapshot) error = %v, want ErrSnapshotMismatch", err)
	}
	report.SnapshotVersion = 3
	report.OwnershipEpoch = 10
	if err := store.ReplaceRuntime(context.Background(), report, time.Minute); !errors.Is(err, ErrSnapshotMismatch) {
		t.Fatalf("ReplaceRuntime(ahead epoch) error = %v, want ErrSnapshotMismatch", err)
	}
}

func TestMemoryStoreExpiredSameIdentitySessionDoesNotReactivateOldReport(t *testing.T) {
	now := time.Date(2026, 7, 30, 15, 0, 0, 0, time.UTC)
	store := newRuntimeStore(t, &now)
	registerRuntimeSession(t, store, "session-a", time.Second)
	if err := store.ReplaceRuntime(context.Background(), Report{
		WorkerID: "worker-a", SessionID: "session-a", Sequence: 1,
		SnapshotVersion: 3, OwnershipEpoch: 9, ObservedAt: now,
	}, time.Minute); err != nil {
		t.Fatalf("ReplaceRuntime(): %v", err)
	}
	session := store.sessions["worker-a"]
	session.expiresAt = now.Add(time.Second)
	store.sessions["worker-a"] = session
	now = now.Add(2 * time.Second)
	registerRuntimeSession(t, store, "session-a", time.Minute)
	got, err := store.ReadRuntime(context.Background(), []OwnedProxy{{
		ProxyID: "proxy-a", WorkerID: "worker-a", OwnershipEpoch: 9,
	}})
	if err != nil || len(got) != 1 || got[0].Fresh {
		t.Fatalf("ReadRuntime(after re-register) = %+v, %v; want stale", got, err)
	}
}

func newRuntimeStore(t *testing.T, now *time.Time) *MemoryStore {
	t.Helper()
	store, err := NewMemoryStore(func() time.Time { return *now })
	if err != nil {
		t.Fatalf("NewMemoryStore(): %v", err)
	}
	return store
}

func registerRuntimeSession(t *testing.T, store *MemoryStore, sessionID string, ttl time.Duration) {
	t.Helper()
	if err := store.ReplaceSession(context.Background(), Session{
		WorkerID: "worker-a", InstanceID: "instance-a", SessionID: sessionID,
		AckedSnapshotVersion: 3, AckedOwnershipEpoch: 9,
	}, ttl); err != nil {
		t.Fatalf("ReplaceSession(): %v", err)
	}
}

func openControlSession(t *testing.T, store *MemoryStore, sessionID string, ttl time.Duration) {
	t.Helper()
	if err := store.OpenSession(context.Background(), Session{
		WorkerID: "worker-a", InstanceID: "instance-a", SessionID: sessionID, Zone: "zone-a", ProtocolVersion: 1,
		Labels: map[string]string{"region": "test"},
	}, ttl); err != nil {
		t.Fatalf("OpenSession(): %v", err)
	}
}

func controlReference(version, epoch uint64, content string) SnapshotReference {
	return SnapshotReference{
		WorkerID: "worker-a", Version: version, OwnershipEpoch: epoch, Checksum: sha256.Sum256([]byte(content)),
	}
}

func controlReport(now time.Time, sessionID string, sequence, version, epoch uint64) Report {
	return Report{
		WorkerID: "worker-a", SessionID: sessionID, Sequence: sequence, SnapshotVersion: version,
		OwnershipEpoch: epoch, ObservedAt: now, Counters: []Counter{{ProxyID: "proxy-a", Active: 1}},
	}
}