package main

import (
	"context"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"syscall"
	"time"
)

var hostname string

func init() {
	var err error
	hostname, err = os.Hostname()
	if err != nil {
		hostname = "unknown"
	}
	if h := os.Getenv("HOSTNAME"); h != "" {
		hostname = h
	}
}

type statusRecorder struct {
	http.ResponseWriter
	status int
}

func (r *statusRecorder) WriteHeader(code int) {
	r.status = code
	r.ResponseWriter.WriteHeader(code)
}

// Forward http.Flusher so chunked-streaming handlers can flush via *statusRecorder.
// Embedded http.ResponseWriter does not promote Flush() — that method is on the
// concrete net/http response type, not on the ResponseWriter interface.
func (r *statusRecorder) Flush() {
	if f, ok := r.ResponseWriter.(http.Flusher); ok {
		f.Flush()
	}
}

func loggingMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		start := time.Now()
		rec := &statusRecorder{ResponseWriter: w, status: 200}
		next.ServeHTTP(rec, r)
		path := r.URL.RequestURI()
		log.Printf("ts=%s method=%s path=%s remote=%s hostname=%s status=%d duration_ms=%d",
			start.UTC().Format("2006-01-02T15:04:05Z"),
			r.Method,
			path,
			r.RemoteAddr,
			hostname,
			rec.status,
			time.Since(start).Milliseconds(),
		)
	})
}

func handleFast(w http.ResponseWriter, r *http.Request) {
	fmt.Fprint(w, "fast ok\n")
}

func handleSleep(w http.ResponseWriter, r *http.Request) {
	n := 1
	if s := r.URL.Query().Get("seconds"); s != "" {
		v, err := strconv.Atoi(s)
		if err == nil && v >= 0 && v <= 600 {
			n = v
		}
	}
	log.Printf("ts=%s hostname=%s path=/sleep seconds=%d remote=%s msg=start",
		time.Now().UTC().Format("2006-01-02T15:04:05Z"), hostname, n, r.RemoteAddr)

	select {
	case <-time.After(time.Duration(n) * time.Second):
		fmt.Fprintf(w, "slept %ds\n", n)
	case <-r.Context().Done():
		log.Printf("ts=%s hostname=%s path=/sleep seconds=%d remote=%s msg=client_disconnected",
			time.Now().UTC().Format("2006-01-02T15:04:05Z"), hostname, n, r.RemoteAddr)
	}
}

func handleStream(w http.ResponseWriter, r *http.Request) {
	totalSec := 10
	intervalSec := 1.0

	if s := r.URL.Query().Get("seconds"); s != "" {
		v, err := strconv.Atoi(s)
		if err == nil && v > 0 && v <= 600 {
			totalSec = v
		}
	}
	if s := r.URL.Query().Get("interval"); s != "" {
		v, err := strconv.ParseFloat(s, 64)
		if err == nil && v > 0 {
			intervalSec = v
		}
	}

	flusher, ok := w.(http.Flusher)
	if !ok {
		http.Error(w, "streaming unsupported", http.StatusInternalServerError)
		return
	}

	// Set chunked transfer before writing first byte.
	w.Header().Set("Transfer-Encoding", "chunked")
	w.Header().Set("Content-Type", "text/plain")
	w.WriteHeader(http.StatusOK)

	start := time.Now()
	deadline := start.Add(time.Duration(totalSec) * time.Second)
	interval := time.Duration(float64(time.Second) * intervalSec)
	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	log.Printf("ts=%s hostname=%s path=/stream seconds=%d interval=%.2f remote=%s msg=start",
		start.UTC().Format("2006-01-02T15:04:05Z"), hostname, totalSec, intervalSec, r.RemoteAddr)

	i := 0
	for {
		select {
		case t := <-ticker.C:
			if t.After(deadline) {
				return
			}
			ms := t.Sub(start).Milliseconds()
			fmt.Fprintf(w, "chunk %d @ %dms\n", i, ms)
			flusher.Flush()
			log.Printf("ts=%s hostname=%s path=/stream chunk=%d ms=%d remote=%s msg=chunk",
				t.UTC().Format("2006-01-02T15:04:05Z"), hostname, i, ms, r.RemoteAddr)
			i++
		case <-r.Context().Done():
			log.Printf("ts=%s hostname=%s path=/stream chunks=%d remote=%s msg=client_disconnected",
				time.Now().UTC().Format("2006-01-02T15:04:05Z"), hostname, i, r.RemoteAddr)
			return
		}
	}
}

func handleHealthz(w http.ResponseWriter, r *http.Request) {
	fmt.Fprint(w, "ok\n")
}

func handleRoot(w http.ResponseWriter, r *http.Request) {
	fmt.Fprint(w, "service-a-backend\n")
}

func main() {
	port := "8080"
	if p := os.Getenv("PORT"); p != "" {
		port = p
	}

	mux := http.NewServeMux()
	mux.HandleFunc("/fast", handleFast)
	mux.HandleFunc("/sleep", handleSleep)
	mux.HandleFunc("/stream", handleStream)
	mux.HandleFunc("/healthz", handleHealthz)
	mux.HandleFunc("/", handleRoot)

	srv := &http.Server{
		Addr:    ":" + port,
		Handler: loggingMiddleware(mux),
	}

	quit := make(chan os.Signal, 1)
	signal.Notify(quit, syscall.SIGTERM, syscall.SIGINT)

	go func() {
		log.Printf("ts=%s hostname=%s msg=listening port=%s",
			time.Now().UTC().Format("2006-01-02T15:04:05Z"), hostname, port)
		if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			log.Fatalf("listen: %v", err)
		}
	}()

	<-quit
	log.Printf("ts=%s hostname=%s msg=received_SIGTERM_shutting_down",
		time.Now().UTC().Format("2006-01-02T15:04:05Z"), hostname)

	// 300s gives in-flight /sleep?seconds=600 requests time to complete.
	ctx, cancel := context.WithTimeout(context.Background(), 300*time.Second)
	defer cancel()
	if err := srv.Shutdown(ctx); err != nil {
		log.Printf("ts=%s hostname=%s msg=shutdown_error err=%v",
			time.Now().UTC().Format("2006-01-02T15:04:05Z"), hostname, err)
	}
	log.Printf("ts=%s hostname=%s msg=shutdown_complete",
		time.Now().UTC().Format("2006-01-02T15:04:05Z"), hostname)
}