package proc

import (
	"context"
	"log/slog"
	"os"
	"os/signal"
	"sync"
	"syscall"
)

// RunWithSignals runs fn under a context that is cancelled on SIGINT/SIGTERM.
// The moment a signal arrives, it logs an acknowledgment so the user gets
// immediate feedback that shutdown has started (instead of seeming to hang
// while child processes drain). On signal or fn returning, onShutdown (if
// non-nil) is invoked exactly once. A second signal aborts the shutdown and
// force-exits the process.
func RunWithSignals(parent context.Context, fn func(context.Context) error, onShutdown func()) error {
	return RunWithSignalsAndForceExit(parent, fn, onShutdown, nil)
}

// RunWithSignalsAndForceExit behaves like RunWithSignals but also invokes
// onForceExit (if non-nil) before force-exiting on a second signal. Use it
// when child processes live in their own process groups and would otherwise
// survive the parent's os.Exit — typical for dev tooling that supervises a
// detached daemon or Electron host.
func RunWithSignalsAndForceExit(
	parent context.Context,
	fn func(context.Context) error,
	onShutdown func(),
	onForceExit func(),
) error {
	signals := make(chan os.Signal, 2)
	signal.Notify(signals, os.Interrupt, syscall.SIGTERM)
	defer signal.Stop(signals)
	return runWithSignalChannel(parent, fn, onShutdown, signals, func(code int) {
		if onForceExit != nil {
			onForceExit()
		}
		os.Exit(code)
	})
}

func runWithSignalChannel(
	parent context.Context,
	fn func(context.Context) error,
	onShutdown func(),
	signals <-chan os.Signal,
	forceExit func(int),
) error {
	ctx, cancel := context.WithCancel(parent)
	defer cancel()

	runDone := make(chan error, 1)
	go func() {
		runDone <- fn(ctx)
	}()

	var (
		runErr      error
		cleanupDone = make(chan struct{})
		cleanupOnce sync.Once
		runFinished bool
		cleanedUp   bool
		signalCount int
		parentDone  = parent.Done()
	)
	startCleanup := func() {
		cleanupOnce.Do(func() {
			go func() {
				if onShutdown != nil {
					onShutdown()
				}
				close(cleanupDone)
			}()
		})
	}
	for {
		if runFinished && cleanedUp {
			return runErr
		}
		select {
		case <-parentDone:
			parentDone = nil
			cancel()
			startCleanup()
		case <-signals:
			signalCount++
			if signalCount == 1 {
				slog.Info("shutdown signal received, cleaning up...")
				cancel()
				startCleanup()
				continue
			}
			slog.Error("second signal received; forcing exit")
			forceExit(130)
		case runErr = <-runDone:
			runFinished = true
			startCleanup()
		case <-cleanupDone:
			cleanedUp = true
			cleanupDone = nil
		}
	}
}