package proc

import (
	"context"
	"os"
	"sync/atomic"
	"testing"
	"time"
)

func TestRunWithSignalChannelFirstSignalCancelsAndCleansUp(t *testing.T) {
	signals := make(chan os.Signal, 2)
	canceled := make(chan struct{})
	cleanupStarted := make(chan struct{})
	releaseCleanup := make(chan struct{})
	forced := make(chan int, 1)
	result := make(chan error, 1)

	go func() {
		result <- runWithSignalChannel(context.Background(), func(ctx context.Context) error {
			<-ctx.Done()
			close(canceled)
			return nil
		}, func() {
			close(cleanupStarted)
			<-releaseCleanup
		}, signals, func(code int) {
			forced <- code
		})
	}()

	signals <- os.Interrupt
	waitForSignalTest(t, canceled, "context cancellation")
	waitForSignalTest(t, cleanupStarted, "cleanup")
	select {
	case code := <-forced:
		t.Fatalf("first signal forced exit with code %d", code)
	case <-time.After(30 * time.Millisecond):
	}

	close(releaseCleanup)
	waitForSignalResult(t, result)
}

func TestRunWithSignalChannelSecondSignalForcesExit(t *testing.T) {
	signals := make(chan os.Signal, 2)
	cleanupStarted := make(chan struct{})
	releaseCleanup := make(chan struct{})
	forced := make(chan int, 1)
	result := make(chan error, 1)

	go func() {
		result <- runWithSignalChannel(context.Background(), func(ctx context.Context) error {
			<-ctx.Done()
			return nil
		}, func() {
			close(cleanupStarted)
			<-releaseCleanup
		}, signals, func(code int) {
			forced <- code
		})
	}()

	signals <- os.Interrupt
	waitForSignalTest(t, cleanupStarted, "cleanup")
	signals <- os.Interrupt

	select {
	case code := <-forced:
		if code != 130 {
			t.Fatalf("force exit code = %d, want 130", code)
		}
	case <-time.After(time.Second):
		t.Fatal("second signal did not force exit")
	}

	close(releaseCleanup)
	waitForSignalResult(t, result)
}

func TestRunWithSignalsAndForceExitInvokesHandlerBeforeOsExit(t *testing.T) {
	signals := make(chan os.Signal, 2)
	cleanupStarted := make(chan struct{})
	releaseCleanup := make(chan struct{})
	forceHandlerCalled := make(chan struct{}, 1)
	forced := make(chan int, 1)

	done := make(chan error, 1)
	go func() {
		done <- runWithSignalChannel(context.Background(), func(ctx context.Context) error {
			<-ctx.Done()
			return nil
		}, func() {
			close(cleanupStarted)
			<-releaseCleanup
		}, signals, func(code int) {
			select {
			case forceHandlerCalled <- struct{}{}:
			default:
			}
			forced <- code
		})
	}()

	signals <- os.Interrupt
	waitForSignalTest(t, cleanupStarted, "cleanup")
	signals <- os.Interrupt

	waitForSignalTest(t, forceHandlerCalled, "force handler before exit")
	select {
	case code := <-forced:
		if code != 130 {
			t.Fatalf("force exit code = %d, want 130", code)
		}
	case <-time.After(time.Second):
		t.Fatal("second signal did not force exit")
	}

	close(releaseCleanup)
	waitForSignalResult(t, done)
}

func TestRunWithSignalChannelNormalReturnCleansUpOnce(t *testing.T) {
	var cleanupCalls atomic.Int32
	forced := make(chan int, 1)

	err := runWithSignalChannel(context.Background(), func(context.Context) error {
		return nil
	}, func() {
		cleanupCalls.Add(1)
	}, make(chan os.Signal), func(int) {
		forced <- 1
	})
	if err != nil {
		t.Fatalf("runWithSignalChannel() error = %v", err)
	}
	select {
	case <-forced:
		t.Fatal("force exit called")
	default:
	}
	if got := cleanupCalls.Load(); got != 1 {
		t.Fatalf("cleanup calls = %d, want 1", got)
	}
}

func TestRunWithSignalChannelParentCancellationCleansUp(t *testing.T) {
	parent, cancel := context.WithCancel(context.Background())
	cleanupCalled := make(chan struct{})
	forced := make(chan int, 1)
	result := make(chan error, 1)

	go func() {
		result <- runWithSignalChannel(parent, func(ctx context.Context) error {
			<-ctx.Done()
			return ctx.Err()
		}, func() {
			close(cleanupCalled)
		}, make(chan os.Signal), func(code int) {
			forced <- code
		})
	}()

	cancel()
	waitForSignalTest(t, cleanupCalled, "cleanup")
	select {
	case err := <-result:
		if err != context.Canceled {
			t.Fatalf("runWithSignalChannel() error = %v, want context.Canceled", err)
		}
	case <-time.After(time.Second):
		t.Fatal("timed out waiting for RunWithSignals result")
	}
	select {
	case code := <-forced:
		t.Fatalf("force exit called with code %d", code)
	default:
	}
}

func waitForSignalTest(t *testing.T, ch <-chan struct{}, what string) {
	t.Helper()
	select {
	case <-ch:
	case <-time.After(time.Second):
		t.Fatalf("timed out waiting for %s", what)
	}
}

func waitForSignalResult(t *testing.T, result <-chan error) {
	t.Helper()
	select {
	case err := <-result:
		if err != nil {
			t.Fatalf("runWithSignalChannel() error = %v", err)
		}
	case <-time.After(time.Second):
		t.Fatal("timed out waiting for RunWithSignals result")
	}
}