package discovery

import (
	"sync"
	"testing"
	"time"

	me "github.com/tech/sendico/pkg/messaging/envelope"
	"go.uber.org/zap"
)

type recordingProducer struct {
	mu     sync.Mutex
	counts map[string]int
}

func newRecordingProducer() *recordingProducer {
	return &recordingProducer{
		counts: make(map[string]int),
	}
}

func (p *recordingProducer) SendMessage(env me.Envelope) error {
	if p == nil || env == nil {
		return nil
	}
	sig := env.GetSignature()
	if sig == nil {
		return nil
	}

	p.mu.Lock()
	p.counts[sig.ToString()]++
	p.mu.Unlock()
	return nil
}

func (p *recordingProducer) count(event string) int {
	if p == nil {
		return 0
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	return p.counts[event]
}

func waitForAtLeast(t *testing.T, timeout time.Duration, countFn func() int, want int, label string) {
	t.Helper()

	deadline := time.Now().Add(timeout)
	for time.Now().Before(deadline) {
		if countFn() >= want {
			return
		}
		time.Sleep(20 * time.Millisecond)
	}
	t.Fatalf("timed out waiting for %s: got=%d want_at_least=%d", label, countFn(), want)
}

func TestHeartbeatIntervalForHealth_Default(t *testing.T) {
	got := heartbeatIntervalForHealth(HealthParams{})
	want := time.Duration(DefaultHealthIntervalSec) * time.Second
	if got != want {
		t.Fatalf("unexpected heartbeat interval: got=%s want=%s", got, want)
	}
}

func TestReannounceIntervalForHealth_Default(t *testing.T) {
	got := heartbeatIntervalForHealth(HealthParams{}) * defaultReannounceHeartbeatFactor
	want := time.Duration(DefaultHealthIntervalSec*defaultReannounceHeartbeatFactor) * time.Second
	if got != want {
		t.Fatalf("unexpected reannounce interval: got=%s want=%s", got, want)
	}
}

func TestAnnouncerPeriodicReannounce(t *testing.T) {
	producer := newRecordingProducer()
	announce := Announcement{
		Service:    "FX_INGESTOR",
		Operations: []string{OperationFXIngest},
		Health: HealthParams{
			IntervalSec: 1,
			TimeoutSec:  2,
		},
	}

	announcer := NewAnnouncer(zap.NewNop(), producer, "fx_ingestor", announce)
	announcer.Start()
	defer announcer.Stop()

	announceEvent := ServiceAnnounceEvent().ToString()
	heartbeatEvent := HeartbeatEvent().ToString()

	waitForAtLeast(t, 1*time.Second, func() int { return producer.count(announceEvent) }, 1, "initial announce")
	waitForAtLeast(t, 1*time.Second, func() int { return producer.count(heartbeatEvent) }, 1, "initial heartbeat")

	// With 1s heartbeat and factor=N, periodic re-announce should happen in ~N seconds.
	timeout := time.Duration(defaultReannounceHeartbeatFactor+2) * time.Second
	waitForAtLeast(t, timeout, func() int { return producer.count(announceEvent) }, 2, "periodic re-announce")
}