package main

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"sync"
	"sync/atomic"
	"time"
)

type App struct {
	ready                            atomic.Bool
	started, exited, release         chan struct{}
	startOnce, exitOnce, releaseOnce sync.Once
}

func NewApp() *App {
	a := &App{started: make(chan struct{}), exited: make(chan struct{}), release: make(chan struct{})}
	a.ready.Store(true)
	return a
}
func (a *App) Release() { a.releaseOnce.Do(func() { close(a.release) }) }
func (a *App) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	switch r.URL.Path {
	case "/live":
		w.WriteHeader(200)
	case "/ready":
		if !a.ready.Load() {
			http.Error(w, "draining", 503)
			return
		}
		w.WriteHeader(200)
	case "/work":
		if !a.ready.Load() {
			http.Error(w, "draining", 503)
			return
		}
		a.startOnce.Do(func() { close(a.started) })
		defer a.exitOnce.Do(func() { close(a.exited) })
		select {
		case <-a.release:
			fmt.Fprint(w, "done")
		case <-r.Context().Done():
			return
		}
	default:
		http.NotFound(w, r)
	}
}

type Running struct {
	App     *App
	Server  *http.Server
	URL     string
	Stopped chan error
}

func Start() (*Running, error) {
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		return nil, err
	}
	a := NewApp()
	srv := &http.Server{Handler: a, ReadHeaderTimeout: time.Second}
	r := &Running{a, srv, "http://" + ln.Addr().String(), make(chan error, 1)}
	go func() { r.Stopped <- srv.Serve(ln) }()
	return r, nil
}
func Get(c *http.Client, url string) (int, error) {
	resp, err := c.Get(url)
	if err != nil {
		return 0, err
	}
	defer resp.Body.Close()
	_, err = io.Copy(io.Discard, resp.Body)
	return resp.StatusCode, err
}
func run() error {
	r, err := Start()
	if err != nil {
		return err
	}
	defer r.Server.Close()
	defer r.App.Release()
	c := &http.Client{Timeout: 2 * time.Second}
	defer c.CloseIdleConnections()
	status, err := Get(c, r.URL+"/ready")
	if err != nil {
		return err
	}
	fmt.Println("ready before:", status)
	work := make(chan int, 1)
	errs := make(chan error, 1)
	go func() { code, e := Get(c, r.URL+"/work"); work <- code; errs <- e }()
	<-r.App.started
	r.App.ready.Store(false)
	status, err = Get(c, r.URL+"/ready")
	if err != nil {
		return err
	}
	fmt.Println("ready while draining:", status)
	status, err = Get(c, r.URL+"/live")
	if err != nil {
		return err
	}
	fmt.Println("live while draining:", status)
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	done := make(chan error, 1)
	go func() { done <- r.Server.Shutdown(ctx) }()
	if err = <-r.Stopped; !errors.Is(err, http.ErrServerClosed) {
		return err
	}
	select {
	case err = <-done:
		return fmt.Errorf("shutdown returned before request completed: %v", err)
	default:
		fmt.Println("shutdown waits for active request:", true)
	}
	r.App.Release()
	fmt.Println("active request status:", <-work)
	if err = <-errs; err != nil {
		return err
	}
	err = <-done
	fmt.Println("shutdown completed:", err == nil)
	return err
}
func main() {
	if err := run(); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}