package server

import (
	"log/slog"
	"net/http"
	"time"

	"g.fc.run/ai/byre/internal/trace"
)

const slowRequest = 200 * time.Millisecond

type timedWriter struct {
	http.ResponseWriter
	start  time.Time
	trace  *trace.Trace
	status int
}

func (w *timedWriter) WriteHeader(status int) {
	if w.status == 0 {
		w.status = status
		w.Header().Set("Server-Timing", w.trace.Header(time.Since(w.start)))
	}
	w.ResponseWriter.WriteHeader(status)
}

func (w *timedWriter) Write(b []byte) (int, error) {
	if w.status == 0 {
		w.WriteHeader(http.StatusOK)
	}
	return w.ResponseWriter.Write(b)
}

func (w *timedWriter) Unwrap() http.ResponseWriter {
	return w.ResponseWriter
}

func timing(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ctx, t := trace.New(r.Context())
		tw := &timedWriter{ResponseWriter: w, start: time.Now(), trace: t}
		next.ServeHTTP(tw, r.WithContext(ctx))
		took := time.Since(tw.start)
		if took < slowRequest || tw.status == http.StatusSwitchingProtocols || tw.Header().Get("Content-Type") == "text/event-stream" {
			return
		}
		slog.Info("request slow", "method", r.Method, "path", r.URL.Path, "status", tw.status, "took", took, "spans", t.String())
	})
}