package main

import (
	"context"
	"fmt"
	"net"
	"net/http"
	"os"
	"os/exec"
	"os/signal"
	"sync"
	"syscall"
	"time"
)

// tcpKeepAliveListener sets TCP keep-alive timeouts on accepted connections.
// It's used by ListenAndServe and ListenAndServeTLS so dead TCP connections
// (e.g. closing laptop mid-download) eventually go away. This is code from
// net/http/server.go.
type tcpKeepAliveListener struct {
	*net.TCPListener
}

// Accept accepts a TCP connection while setting keep-alive timeouts.
func (ln tcpKeepAliveListener) Accept() (c net.Conn, err error) {
	tc, err := ln.AcceptTCP()
	if err != nil {
		return
	}
	tc.SetKeepAlive(true)
	tc.SetKeepAlivePeriod(3 * time.Minute)
	return tc, nil
}

// Serve starts an HTTP server using the DefaultServeMux. This function does
// not return until the server has been stopped and all requests have finished.
//
// The server is controlled using Unix signals:
//
//   kill -s SIGINT <pid> (or Ctrl-C) stops the server, waiting for any active
//   requests to finish.
//
//   kill -s SIGUSR1 <pid> restarts the server with no downtime, using its (new)
//   binary. Current requests will finish before the old server exits and while
//   the new server already processes new requests.
//
//   kill -s SIGTERM <pid> stops the server without waiting for active requests
//   to finish.
func Serve(address string) error {
	// The general idea for this is described here:
	// https://grisha.org/blog/2014/06/03/graceful-restart-in-golang/

	// We need to shut down gracefully when the user hits Ctrl-C.
	sig := make(chan os.Signal, 1)
	signal.Notify(sig, syscall.SIGINT, syscall.SIGUSR1, syscall.SIGTERM)

	// Create a server.
	server := &http.Server{Addr: address}

	// Start the server.
	var (
		listener *net.TCPListener
		wg       sync.WaitGroup
	)
	wg.Add(1)
	go func() {
		defer wg.Done()

		// If this is a forked child process, we'll use its connection.
		isFork := os.Getenv("FORKED_SERVER") != ""

		var (
			ln  net.Listener
			err error
		)
		if isFork {
			// It's a fork. Get the file that was handed over.
			fmt.Printf("%d Getting existing listener for %s\n", pid, address)
			file := os.NewFile(3, "")
			ln, err = net.FileListener(file)
			if err != nil {
				fmt.Printf("%d Cannot use existing listener: %s\n", pid, err)
				sig <- syscall.SIGTERM
				return
			}

			// Tell the parent to stop the server now.
			parent := syscall.Getppid()
			fmt.Printf("%d Telling parent process (%d) to stop server\n", pid, parent)
			syscall.Kill(parent, syscall.SIGTERM)

			// Give the parent some time.
			time.Sleep(100 * time.Millisecond)
		} else {
			// It's a new server.
			fmt.Printf("%d Starting web server on %s\n", pid, address)
			ln, err = net.Listen("tcp", address)
			if err != nil {
				fmt.Printf("%d Cannot listen to %s: %s\n", pid, address, err)
				sig <- syscall.SIGTERM
				return
			}
		}

		// We can start the server now.
		fmt.Println(pid, "Serving requests...")
		listener = ln.(*net.TCPListener)
		if err = server.Serve(tcpKeepAliveListener{listener}); err != nil {
			fmt.Printf("%d Web server was shut down: %s\n", pid, err)
		}
		fmt.Println(pid, "Web server has finished")
	}()

	// Wait for the interrupt signal.
	s := <-sig
	switch s {
	case syscall.SIGTERM:
		// Go for the program exit. Don't wait for the server to finish.
		fmt.Println(pid, "Received SIGTERM, exiting without waiting for the web server to shut down")
		return nil
	case syscall.SIGINT:
		// Stop the server gracefully.
		fmt.Println(pid, "Received SIGINT")
	case syscall.SIGUSR1:
		// Spawn a child process.
		fmt.Println(pid, "Received SIGUSR1")
		var args []string
		if len(os.Args) > 1 {
			args = os.Args[1:]
		}
		file, err := listener.File()
		if err != nil {
			fmt.Printf("%d Listener did not return file, not forking: %s\n", pid, err)
		} else {
			cmd := exec.Command(os.Args[0], args...)
			cmd.Stdout = os.Stdout
			cmd.Stderr = os.Stderr
			cmd.ExtraFiles = []*os.File{file}
			cmd.Env = append(os.Environ(), "FORKED_SERVER=1")
			if err := cmd.Start(); err != nil {
				fmt.Printf("%d Fork did not succeed: %s\n", pid, err)
			}
			fmt.Printf("%d Started child process %d, waiting for its ready signal\n", pid, cmd.Process.Pid)

			// We have a child process. A SIGTERM means the child process is ready to
			// start its server.
			<-sig
		}
	}

	// Force the server to shut down.
	fmt.Println(pid, "Shutting down web server and waiting for requests to finish...")
	defer fmt.Println(pid, "Requests have finished")
	if err := server.Shutdown(context.Background()); err != nil {
		return fmt.Errorf("Shutdown failed: %s", err)
	}
	wg.Wait()

	return nil
}