// Copyright (c) 2026 Storm Dragon
// SPDX-License-Identifier: GPL-3.0-or-later

package tentacle

import (
	"context"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"syscall"
	"time"
)

var ErrSupervision = errors.New("tentacle process supervision failed")

// RunSupervisor is a private per-session helper. Descriptors 0..2 carry only
// the Tentacle protocol and diagnostics; 3 is a reader-owned lifetime pipe and
// 4 is a bounded Spec. The child never inherits either private descriptor.
// A separate subreaper keeps ownership of double-forked/setsid descendants
// without making the reader a subreaper for unrelated desktop helpers.
func RunSupervisor(ctx context.Context) error {
	// ExtraFiles arrive in blocking mode. Register the private pipes with Go's
	// poller so Close interrupts their readers on cancellation and child exit.
	for _, descriptor := range []int{3, 4} {
		syscall.CloseOnExec(descriptor)
		if err := syscall.SetNonblock(descriptor, true); err != nil {
			return fmt.Errorf("%w: private pipe: %w", ErrSupervision, err)
		}
	}
	lifetime := os.NewFile(3, "tentacle-lifetime")
	configuration := os.NewFile(4, "tentacle-spec")
	if lifetime == nil || configuration == nil {
		return ErrSupervision
	}
	defer configuration.Close()
	// PR_SET_CHILD_SUBREAPER is a Linux facility independent of the init system.
	if _, _, errno := syscall.Syscall6(syscall.SYS_PRCTL, 36, 1, 0, 0, 0, 0); errno != 0 {
		lifetime.Close()
		return fmt.Errorf("%w: subreaper: %w", ErrSupervision, errno)
	}
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	readerDone := make(chan struct{})
	go func() {
		defer close(readerDone)
		_, _ = io.Copy(io.Discard, lifetime)
		cancel()
	}()
	defer func() { lifetime.Close(); <-readerDone }()
	// Closing this descriptor also interrupts a partial startup on termination.
	stop := context.AfterFunc(ctx, func() { configuration.Close() })
	defer stop()
	data, err := io.ReadAll(io.LimitReader(configuration, MaxFrameBytes+1))
	if err != nil || len(data) > MaxFrameBytes {
		// No child exists yet. A partial configuration can result from startup
		// cancellation. Protocol EOF fails registration, while a zero helper
		// exit certifies cleanup and lets the host release its admission slot.
		return nil
	}
	configuration.Close()
	var spec Spec
	if decode_strict(data, &spec) != nil || validate_spec(spec) != nil {
		return nil
	}
	if ctx.Err() != nil {
		return nil
	}
	// Pdeathsig follows the creating thread; retain it for the child lifetime.
	runtime.LockOSThread()
	defer runtime.UnlockOSThread()
	command := exec.Command(spec.Program[0], spec.Program[1:]...)
	command.Dir = spec.Directory
	command.Stdin, command.Stdout, command.Stderr = os.Stdin, os.Stdout, os.Stderr
	command.SysProcAttr = &syscall.SysProcAttr{Setpgid: true, Pdeathsig: syscall.SIGKILL}
	childSignals := make(chan os.Signal, 1)
	signal.Notify(childSignals, syscall.SIGCHLD)
	defer signal.Stop(childSignals)
	if err := command.Start(); err != nil {
		// No process was started. Protocol EOF reports the startup failure;
		// successful cleanup keeps the manager available for a later retry.
		return nil
	}
	exited := make(chan struct{})
	go func() { _ = command.Wait(); close(exited) }()
	var supervisionError error
running:
	for {
		select {
		case <-ctx.Done():
			break running
		case <-exited:
			break running
		case <-childSignals:
			// The subreaper owns orphaned descendants for the entire session,
			// not just teardown. Only command.Wait may reap the main child.
			supervisionError = reap_exited_supervised_children(command.Process.Pid)
			if supervisionError != nil {
				break running
			}
		}
	}
	_ = command.Process.Kill()
	<-exited
	return errors.Join(supervisionError, reap_supervised_children(DefaultResourceLimits().CleanupTimeout))
}

func supervised_children() ([]int, error) {
	paths, err := filepath.Glob("/proc/self/task/*/children")
	if err != nil || len(paths) == 0 {
		return nil, ErrSupervision
	}
	var children []int
	for _, path := range paths {
		data, err := os.ReadFile(path)
		if os.IsNotExist(err) || errors.Is(err, syscall.ESRCH) { // A runtime thread exited during the scan.
			continue
		}
		if err != nil {
			return nil, ErrSupervision
		}
		for _, value := range strings.Fields(string(data)) {
			pid, err := strconv.Atoi(value)
			if err != nil || pid <= 1 {
				return nil, ErrSupervision
			}
			children = append(children, pid)
		}
	}
	return children, nil
}

func reap_exited_supervised_children(mainPID int) error {
	children, err := supervised_children()
	if err != nil {
		return err
	}
	for _, childPID := range children {
		if childPID == mainPID {
			continue
		}
		for {
			_, err := syscall.Wait4(childPID, nil, syscall.WNOHANG, nil)
			if err == syscall.EINTR {
				continue
			}
			if err != nil && err != syscall.ECHILD {
				return fmt.Errorf("%w: reap exited descendant: %w", ErrSupervision, err)
			}
			break
		}
	}
	return nil
}

// Only this helper's children are killed. A child PID cannot be reused until
// it has been reaped, and no other goroutine calls Wait after command.Wait joins.
// Killing direct children adopts the next generation, including detached ones.
func reap_supervised_children(timeout time.Duration) error {
	deadline := time.NewTimer(timeout)
	defer deadline.Stop()
	tick := time.NewTicker(5 * time.Millisecond)
	defer tick.Stop()
	for {
		children, err := supervised_children()
		if err != nil {
			return err
		}
		for _, pid := range children {
			if err := syscall.Kill(pid, syscall.SIGKILL); err != nil && err != syscall.ESRCH {
				return fmt.Errorf("%w: terminate descendant: %w", ErrSupervision, err)
			}
		}
		for {
			var status syscall.WaitStatus
			pid, err := syscall.Wait4(-1, &status, syscall.WNOHANG, nil)
			if err == syscall.ECHILD {
				return nil
			}
			if err == syscall.EINTR {
				continue
			}
			if err != nil {
				return fmt.Errorf("%w: reap descendant: %w", ErrSupervision, err)
			}
			if pid == 0 {
				break
			}
		}
		select {
		case <-deadline.C:
			return fmt.Errorf("%w: descendant cleanup deadline expired", ErrSupervision)
		case <-tick.C:
		}
	}
}