package main

import (
	"errors"
	"flag"
	"io"
	"log"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"syscall"
)

var (
	c       = flag.Int("c", 10, "concurrency level, default 10")
	prefork = flag.Bool("prefork", false, "use prefork mode")
	child   = flag.Bool("child", false, "is child process")
)

func main() {
	flag.Parse()

	var ln net.Listener
	var err error

	if *prefork { // 如果要启动子进程模式
		ln = doPrefork(*c)
	} else {
		ln, err = net.Listen("tcp", ":8972")
		if err != nil {
			panic(err)
		}
	}

	start(ln) // 处理 net.Listener
}

func start(ln net.Listener) {
	log.Println("started")
	for {
		conn, e := ln.Accept()
		if e != nil {
			var ne net.Error
			if errors.As(e, &ne) && ne.Temporary() {
				log.Printf("accept temp err: %v", ne)
				continue
			}

			log.Printf("accept err: %v", e)
			return
		}

		go func() {
			_, err := io.Copy(conn, conn) // 实现 echo 协议，将收到的东西原样返回
			if err != nil {
				log.Printf("copy err: %v", err)
			}
		}()
	}
}

func doPrefork(c int) net.Listener {
	if *child {
		// 子进程：从文件描述符恢复监听器
		listener, err := net.FileListener(os.NewFile(3, ""))
		if err != nil {
			log.Fatal(err)
		}

		// 优雅关闭处理
		sigCh := make(chan os.Signal, 1)
		signal.Notify(sigCh, syscall.SIGTERM)
		go func() {
			<-sigCh
			err := listener.Close()
			if err != nil {
				return
			}
		}()

		return listener
	}

	// 主进程：创建监听器并将其传递给子进程
	addr, err := net.ResolveTCPAddr("tcp", ":8972")
	if err != nil {
		log.Fatal(err)
	}
	tcpListener, err := net.ListenTCP("tcp", addr)
	if err != nil {
		log.Fatal(err)
	}
	fl, err := tcpListener.File()
	if err != nil {
		log.Fatal(err)
	}
	// 主进程不需要保持监听
	if err := tcpListener.Close(); err != nil {
		log.Fatal(err)
	}

	// 启动子进程并持续监控
	children := make([]*exec.Cmd, c)

	// 处理信号以优雅关闭
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)

	// 启动子进程管理协程
	go func() {
		<-sigCh
		// 收到信号后通知所有子进程优雅关闭
		for _, child := range children {
			if child.Process != nil {
				if err := child.Process.Signal(syscall.SIGTERM); err != nil {
					log.Fatal(err)
				}
			}
		}
		os.Exit(0)
	}()

	// 声明子进程启动函数以便递归调用
	var startChildProcess func(int)
	startChildProcess = func(i int) {
		cmd := exec.Command(os.Args[0], "-prefork", "-child")
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		cmd.ExtraFiles = []*os.File{fl}
		children[i] = cmd

		if err := cmd.Start(); err != nil {
			log.Fatalf("启动子进程 %d 失败: %v", i, err)
		}

		// 在单独的协程中等待子进程结束并重启
		go func() {
			if err := cmd.Wait(); err != nil {
				log.Printf("子进程 %d 异常退出: %v，正在重启...", i, err)
				startChildProcess(i) // 递归重启
			}
		}()
	}
	// 启动和监控子进程
	for i := range children {
		startChildProcess(i)
	}

	// 主进程保持运行，但不处理连接
	select {} // 阻塞主进程
}