package cluster

import (
	"context"
	"encoding/json"
	"testing"
	"time"

	"armature/store"
)

// mockClusterStore implements store.ClusterStore for unit tests.
type mockClusterStore struct {
	nodes          []store.ClusterNode
	heartbeatRows  int64
	heartbeatErr   error
	sweepCalled    bool
	registerCalled bool
	lastStats      json.RawMessage
}

func (m *mockClusterStore) Register(_ context.Context, nodeID, endpoint string) error {
	m.registerCalled = true
	return nil
}

func (m *mockClusterStore) Heartbeat(_ context.Context, nodeID string, stats json.RawMessage) (int64, error) {
	m.lastStats = stats
	return m.heartbeatRows, m.heartbeatErr
}

func (m *mockClusterStore) SweepStale(_ context.Context, _ time.Duration) (int64, error) {
	m.sweepCalled = true
	return 0, nil
}

func (m *mockClusterStore) ListNodes(_ context.Context) ([]store.ClusterNode, error) {
	return m.nodes, nil
}

func (m *mockClusterStore) Deregister(_ context.Context, _ string) error {
	return nil
}

// mockHub implements ConnCounter.
type mockHub struct{ count int }

func (h *mockHub) ConnCount() int { return h.count }

func TestCollectStats(t *testing.T) {
	hub := &mockHub{count: 5}
	reg := NewRegistry("test-node", "http://localhost:8080", nil, hub, RegistryConfig{
		HeartbeatInterval: 10 * time.Second,
		StaleThreshold:    30 * time.Second,
	})

	data := reg.collectStats()

	var stats map[string]any
	if err := json.Unmarshal(data, &stats); err != nil {
		t.Fatalf("collectStats returned invalid JSON: %v", err)
	}

	// Verify expected keys
	expectedKeys := []string{"goroutines", "heap_alloc", "heap_sys", "gc_cycles", "gc_pause_ns", "uptime_sec", "ws_clients"}
	for _, key := range expectedKeys {
		if _, ok := stats[key]; !ok {
			t.Errorf("missing expected stats key: %s", key)
		}
	}

	// ws_clients should reflect hub count
	if ws, ok := stats["ws_clients"].(float64); !ok || int(ws) != 5 {
		t.Errorf("ws_clients = %v, want 5", stats["ws_clients"])
	}
}

func TestCollectStatsFattened(t *testing.T) {
	hub := &mockHub{count: 3}
	reg := NewRegistry("test-node", "", nil, hub, RegistryConfig{
		HeartbeatInterval: 10 * time.Second,
		StaleThreshold:    30 * time.Second,
	})
	reg.SetSandboxStats(func() (uint64, uint64, float64) { return 100, 5, 12.3 })
	reg.SetTriggerFireCount(func() int64 { return 42 })
	reg.SetExtensionCount(func() int { return 8 })

	data := reg.collectStats()

	var stats map[string]any
	if err := json.Unmarshal(data, &stats); err != nil {
		t.Fatalf("collectStats returned invalid JSON: %v", err)
	}

	// New fattened keys
	fatKeys := []string{
		"stack_in_use", "gc_cpu_pct",
		"extensions_loaded", "starlark_exec_total", "starlark_errors_total",
		"starlark_avg_duration_ms", "trigger_fires_total",
	}
	for _, key := range fatKeys {
		if _, ok := stats[key]; !ok {
			t.Errorf("missing fattened stats key: %s", key)
		}
	}

	if v := stats["extensions_loaded"].(float64); int(v) != 8 {
		t.Errorf("extensions_loaded = %v, want 8", v)
	}
	if v := stats["starlark_exec_total"].(float64); int(v) != 100 {
		t.Errorf("starlark_exec_total = %v, want 100", v)
	}
	if v := stats["trigger_fires_total"].(float64); int(v) != 42 {
		t.Errorf("trigger_fires_total = %v, want 42", v)
	}
}

func TestRegistryStartStop(t *testing.T) {
	ms := &mockClusterStore{heartbeatRows: 1}
	hub := &mockHub{count: 0}

	reg := NewRegistry("test-node", "", ms, hub, RegistryConfig{
		HeartbeatInterval: 50 * time.Millisecond,
		StaleThreshold:    150 * time.Millisecond,
	})

	if err := reg.Start(); err != nil {
		t.Fatalf("Start() error: %v", err)
	}
	if !ms.registerCalled {
		t.Error("Register was not called on Start")
	}

	// Let at least one tick fire
	time.Sleep(100 * time.Millisecond)

	reg.Stop()

	if !ms.sweepCalled {
		t.Error("SweepStale was never called during tick")
	}
	if ms.lastStats == nil {
		t.Error("Heartbeat was never called with stats")
	}
}

func TestNodeID(t *testing.T) {
	reg := NewRegistry("my-node-123", "", nil, &mockHub{}, RegistryConfig{})
	if got := reg.NodeID(); got != "my-node-123" {
		t.Errorf("NodeID() = %q, want %q", got, "my-node-123")
	}
}