package main

import (
	"flag"
	"log"
	"os"
	"os/exec"
	"os/signal"
	"sync"
	"syscall"
	"time"
)

var (
	businessProgramType = flag.String("business-program-type", "", "Type of business program (java or python)")
	businessProgramPath = flag.String("business-program-path", "", "Path to the business program file")
	stopRequested       bool
	currentCmd          *exec.Cmd
	mu                  sync.Mutex
)

func startBusinessProcess(programType, programPath string) *exec.Cmd {
	var cmd *exec.Cmd
	switch programType {
	case "java":
		cmd = exec.Command("java", "-jar", programPath)
	case "python":
		cmd = exec.Command("python", programPath)
	default:
		log.Fatalf("Unsupported business program type: %s", programType)
	}
	return cmd
}

// C:\Users\wdd\go\bin\gox.exe -osarch="linux/amd64" -output "build/watchdog"
func main() {
	// 解析命令行参数
	flag.Parse()
	if *businessProgramType == "" || *businessProgramPath == "" {
		log.Fatal("Missing required flags: -business-program-type and -business-program-path must be specified")
	}
	
	// 信号处理
	signalChan := make(chan os.Signal, 1)
	signal.Notify(signalChan, os.Interrupt, syscall.SIGTERM)
	go func() {
		for sig := range signalChan {
			log.Printf("Received signal: %v", sig)
			mu.Lock()
			stopRequested = true
			if currentCmd != nil && currentCmd.Process != nil {
				// 发送 SIGTERM 给业务进程
				if err := currentCmd.Process.Signal(syscall.SIGTERM); err != nil {
					log.Printf("Failed to send SIGTERM to process: %v", err)
				}
				// 等待 10 秒后强制杀死进程
				time.AfterFunc(10*time.Second, func() {
					mu.Lock()
					defer mu.Unlock()
					if currentCmd != nil && currentCmd.Process != nil {
						log.Println("Graceful shutdown timeout, sending SIGKILL")
						currentCmd.Process.Kill()
					}
				})
			}
			mu.Unlock()
		}
	}()
	
	// 主循环
	for {
		mu.Lock()
		if stopRequested {
			mu.Unlock()
			log.Println("Shutting down due to stop request")
			os.Exit(0)
		}
		mu.Unlock()
		
		cmd := startBusinessProcess(*businessProgramType, *businessProgramPath)
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		
		// 启动业务进程
		if err := cmd.Start(); err != nil {
			log.Printf("Failed to start business process: %v", err)
			time.Sleep(5 * time.Second)
			continue
		}
		
		mu.Lock()
		currentCmd = cmd
		mu.Unlock()
		
		// 等待业务进程退出
		err := cmd.Wait()
		mu.Lock()
		currentCmd = nil
		mu.Unlock()
		
		// 程序异常退出，或者重启
		// 定位分析问题
		if err != nil {
			log.Printf("Business process exited with error: %v", err)
		} else {
			log.Println("Business process exited normally")
		}
		
		mu.Lock()
		if stopRequested {
			mu.Unlock()
			log.Println("Shutting down due to stop request")
			os.Exit(0)
		}
		mu.Unlock()
		
		// 等待 5 秒后重启
		log.Println("Restarting business process in 5 seconds...")
		time.Sleep(5 * time.Second)
	}
}
