package broadcast

import (
	"context"
	"testing"
)

// run is a helper to encapsulate the external loop requirement.
func run[T any](ctx context.Context, w *Worker[T], stop chan struct{}) {
	for {
		if err := w.Work(ctx); err != nil {
			break
		}
		select {
		case <-ctx.Done():
			return
		case <-stop:
			return
		default:
		}
	}
}

func TestWorker_Policies(t *testing.T) {
	tests := []struct {
		name           string
		policy         Policy
		bufferSize     int
		sendCount      int
		readCount      int
		expectBlocked  bool
		expectDelivery bool
	}{
		{
			name:           "Block_DeliverAll",
			policy:         Block,
			bufferSize:     10,
			sendCount:      5,
			readCount:      5,
			expectDelivery: true,
		},
		{
			name:           "Drop_LossOnFull",
			policy:         Drop,
			bufferSize:     1,
			sendCount:      5,
			readCount:      1, // Only read one, others should be dropped
			expectDelivery: true,
		},
		{
			name:           "FireForget_EventualDelivery",
			policy:         FireForget,
			bufferSize:     0, // Unbuffered to force goroutine spawn
			sendCount:      1,
			readCount:      1,
			expectDelivery: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			ctx, cancel := context.WithCancel(context.Background())
			defer cancel()

			input := make(chan int)
			out := make(chan int, tt.bufferSize)
			w := New(input, tt.policy)
			w.Output["test"] = out

			stop := make(chan struct{})
			go run(ctx, w, stop)
			defer close(stop)

			// Send messages
			go func() {
				for i := 0; i < tt.sendCount; i++ {
					input <- i
				}
			}()

			// Collect messages
			received := 0
			for i := 0; i < tt.readCount; i++ {
				select {
				case <-out:
					received++
				case <-ctx.Done():
				}
			}

			if received != tt.readCount {
				t.Errorf("policy %v: expected %d received, got %d", tt.policy, tt.readCount, received)
			}
		})
	}
}

func TestWorker_Lifecycle(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	input := make(chan int)
	w := New(input, Block)

	stop := make(chan struct{})
	go run(ctx, w, stop)
	defer close(stop)

	out := make(chan int, 1)
	id := "client-1"

	// 1. Test Add
	w.AddOutput(id, out)
	input <- 100
	if val := <-out; val != 100 {
		t.Errorf("expected 100, got %d", val)
	}

	// 2. Test Remove
	w.RemoveOutput(id)

	// Use a channel to detect if the worker is still alive but ignoring the output
	done := make(chan bool)
	go func() {
		input <- 200
		done <- true
	}()

	select {
	case <-out:
		t.Error("received message after RemoveOutput")
	case <-done:
		// Success: worker processed message but didn't send to removed output
	}
}

func TestWorker_EdgeCases(t *testing.T) {
	t.Run("ClosedInput", func(t *testing.T) {
		input := make(chan int)
		w := New(input, Block)
		close(input)
		if err := w.Work(context.Background()); err != nil {
			t.Errorf("expected nil on closed input, got %v", err)
		}
	})

	t.Run("ContextCanceled", func(t *testing.T) {
		ctx, cancel := context.WithCancel(context.Background())
		cancel()
		w := New(make(chan int), Block)
		if err := w.Work(ctx); err != context.Canceled {
			t.Errorf("expected context.Canceled, got %v", err)
		}
	})
}