package main

import (
	"context"
	"errors"
	"fmt"
	"net"
	"net/http"
	"os"
	"os/signal"
	"runtime"
	"runtime/debug"
	"sync/atomic"
	"syscall"
	"time"
)

func routes(ready *atomic.Bool, work func(http.ResponseWriter, *http.Request)) http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc("GET /healthz", func(w http.ResponseWriter, r *http.Request) { fmt.Fprintln(w, "ok") })
	mux.HandleFunc("GET /readyz", func(w http.ResponseWriter, r *http.Request) {
		if !ready.Load() {
			http.Error(w, "draining", http.StatusServiceUnavailable)
			return
		}
		fmt.Fprintln(w, "ready")
	})
	mux.HandleFunc("GET /work", work)
	return mux
}
func inspect() {
	fmt.Printf("Go=%s NumCPU=%d default=%d\n", runtime.Version(), runtime.NumCPU(), runtime.GOMAXPROCS(0))
	runtime.GOMAXPROCS(1)
	fmt.Println("explicit=", runtime.GOMAXPROCS(0))
	runtime.SetDefaultGOMAXPROCS()
	fmt.Println("restored=", runtime.GOMAXPROCS(0))
}
func memory() {
	previous := debug.SetMemoryLimit(64 << 20)
	defer debug.SetMemoryLimit(previous)
	live := make([]byte, 80<<20)
	for i := 0; i < len(live); i += 4096 {
		live[i] = 1
	}
	runtime.GC()
	var m runtime.MemStats
	runtime.ReadMemStats(&m)
	fmt.Printf("limit=%dMiB live-buffer=%dMiB managed-above-limit=%t\n", debug.SetMemoryLimit(-1)>>20, len(live)>>20, m.Sys-m.HeapReleased > 64<<20)
	runtime.KeepAlive(live)
}
func run() error {
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()
	var ready atomic.Bool
	ready.Store(true)
	srv := &http.Server{ReadHeaderTimeout: 2 * time.Second}
	srv.Handler = routes(&ready, func(w http.ResponseWriter, r *http.Request) {
		fmt.Println("work-started")
		select {
		case <-time.After(time.Second):
			fmt.Fprintln(w, "finished")
		case <-r.Context().Done():
			return
		}
	})
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		return err
	}
	fmt.Println("listening=" + ln.Addr().String())
	result := make(chan error, 1)
	go func() { result <- srv.Serve(ln) }()
	select {
	case err := <-result:
		return err
	case <-ctx.Done():
	}
	ready.Store(false)
	fmt.Println("draining")
	stop()
	// Teaching delay only; a real deployment must measure endpoint/LB propagation.
	time.Sleep(500 * time.Millisecond)
	drain, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	if err := srv.Shutdown(drain); err != nil {
		_ = srv.Close()
		return err
	}
	if err := <-result; !errors.Is(err, http.ErrServerClosed) {
		return err
	}
	fmt.Println("shutdown-complete")
	return nil
}
func main() {
	if len(os.Args) > 1 {
		switch os.Args[1] {
		case "runtime":
			inspect()
			return
		case "memory":
			memory()
			return
		}
	}
	if err := run(); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}