package main

import (
	"context"
	"runtime"
	"sort"
	"testing"
	"time"
)

func gen(vals ...int) <-chan int {
	ch := make(chan int)
	go func() {
		defer close(ch)
		for _, v := range vals {
			ch <- v
		}
	}()
	return ch
}

func TestMergeAllValues(t *testing.T) {
	out := Merge(context.Background(), gen(1, 2, 3), gen(4, 5), gen(6))
	got := []int{}
	done := make(chan struct{})
	go func() {
		for v := range out {
			got = append(got, v)
		}
		close(done)
	}()
	select {
	case <-done:
	case <-time.After(2 * time.Second):
		t.Fatal("output never closed after all inputs closed")
	}
	sort.Ints(got)
	want := []int{1, 2, 3, 4, 5, 6}
	if len(got) != len(want) {
		t.Fatalf("want %v, got %v", want, got)
	}
	for i := range want {
		if got[i] != want[i] {
			t.Fatalf("want %v, got %v", want, got)
		}
	}
}

func TestNoInputs(t *testing.T) {
	out := Merge[int](context.Background())
	select {
	case _, ok := <-out:
		if ok {
			t.Fatal("value from zero inputs")
		}
	case <-time.After(time.Second):
		t.Fatal("output not closed with zero inputs")
	}
}

func TestCancellationReleasesGoroutines(t *testing.T) {
	before := runtime.NumGoroutine()

	ctx, cancel := context.WithCancel(context.Background())
	// Inputs that never close and keep producing.
	mk := func() <-chan int {
		ch := make(chan int)
		go func() {
			for {
				select {
				case ch <- 1:
				case <-ctx.Done():
					return
				}
			}
		}()
		return ch
	}
	out := Merge(ctx, mk(), mk(), mk())

	// Read a few values, then walk away AND cancel — forwarders must not
	// stay blocked on `out <- v`.
	for i := 0; i < 5; i++ {
		select {
		case <-out:
		case <-time.After(2 * time.Second):
			cancel()
			t.Fatal("no values arriving on merged channel")
		}
	}
	cancel()

	deadline := time.After(2 * time.Second)
	for {
		runtime.GC()
		if runtime.NumGoroutine() <= before+2 { // small slack
			return
		}
		select {
		case <-deadline:
			t.Fatalf("goroutine leak after cancel: before=%d now=%d", before, runtime.NumGoroutine())
		case <-time.After(20 * time.Millisecond):
		}
	}
}