package main

import (
	"context"
	"fmt"
	"log/slog"
	"net"
	"net/http"
	"strings"
	"sync/atomic"
	"time"
)

type operationalHandler struct {
	next     http.Handler
	ready    func(context.Context) error
	hosts    map[string]bool
	slots    chan struct{}
	sequence atomic.Uint64
	logger   *slog.Logger
}

type operationalWriter struct {
	http.ResponseWriter
	status int
}

func (w *operationalWriter) WriteHeader(status int) {
	if w.status == 0 {
		w.status = status
		w.ResponseWriter.WriteHeader(status)
	}
}
func (w *operationalWriter) Write(p []byte) (int, error) {
	if w.status == 0 {
		w.WriteHeader(200)
	}
	return w.ResponseWriter.Write(p)
}

func (h *operationalHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	start := time.Now()
	w.Header().Set("X-Request-ID", fmt.Sprintf("%x-%x", start.UnixNano(), h.sequence.Add(1)))
	out := &operationalWriter{ResponseWriter: w}
	defer func() {
		if recover() != nil {
			if out.status == 0 {
				writeJSON(out, 500, map[string]string{"detail": "Internal Server Error"})
			}
			h.logger.Error("handler panic", "request_id", w.Header().Get("X-Request-ID"))
		}
		status := out.status
		if status == 0 {
			status = 200
		}
		h.logger.Info("request", "request_id", w.Header().Get("X-Request-ID"), "method", r.Method, "status", status, "duration_ms", float64(time.Since(start))/1e6)
	}()
	host := r.Host
	if name, _, err := net.SplitHostPort(host); err == nil {
		host = name
	}
	if !h.hosts[strings.ToLower(host)] {
		writeJSON(out, 400, map[string]string{"detail": "Invalid host"})
		return
	}
	// Forwarded and X-Forwarded-* headers are deliberately not trusted.
	if r.URL.Path == "/healthz" || r.URL.Path == "/readyz" {
		if r.Method != "GET" {
			out.Header().Set("Allow", "GET")
			out.WriteHeader(405)
			return
		}
		if r.URL.Path == "/readyz" {
			ctx, cancel := context.WithTimeout(r.Context(), time.Second)
			defer cancel()
			if err := h.ready(ctx); err != nil {
				writeJSON(out, 503, map[string]string{"status": "unavailable"})
				return
			}
		}
		writeJSON(out, 200, map[string]string{"status": "ok"})
		return
	}
	select {
	case h.slots <- struct{}{}:
		defer func() { <-h.slots }()
	default:
		writeJSON(out, 503, map[string]string{"detail": "Busy"})
		return
	}
	const maxBody = 1 << 20
	if r.ContentLength > maxBody {
		writeJSON(out, 413, map[string]string{"detail": "Request body too large"})
		return
	}
	r.Body = http.MaxBytesReader(out, r.Body, maxBody)
	ctx, cancel := context.WithTimeout(r.Context(), 15*time.Second)
	defer cancel()
	h.next.ServeHTTP(out, r.WithContext(ctx))
}