package mcphttp

import "net/http"

const (
	// cacheControl is what every response of an MCP endpoint carries.
	//
	// private and no-store are middleware.SetPrivateCache's pair, unchanged and
	// for its reason: an answer here depends entirely on the credential the
	// request carried and says nothing about it in the URL, and it may be a
	// PRIVATE repository's data. no-cache would not do — it still permits a cache
	// to *store* the body and merely revalidate, which is the thing no-store
	// forbids.
	//
	// no-transform is the SDK's own, kept. The streamable transport writes
	// `no-cache, no-transform` on every response it produces: no-transform
	// protects the SSE framing from an intermediary that would recompress or
	// rechunk it, and there is no reason to drop it. Only no-cache is replaced.
	cacheControl = "private, no-store, no-transform"

	// vary names what an answer here actually depends on: the unified-login
	// cookie and the bearer token. Same two names the rest of the instance varies
	// on, because it is the same statement.
	vary = "Cookie, Authorization"
)

// PrivateCache marks every response an MCP endpoint writes as one no cache may
// keep, and states what it depends on.
//
// It is not middleware.PrivateCache and cannot be. That one sets the headers
// before the handler runs, which is right for a router whose handlers do not
// touch Cache-Control; the SDK's streamable transport sets Cache-Control itself,
// with Set, from inside the handler, so a value written on the way in is
// overwritten on the way out and the response leaves with `no-cache,
// no-transform` and no Vary at all. The headers are therefore written at the
// last moment they still can be: when the status line is committed and every Set
// the handler was going to make has been made.
//
// The result is not a replacement of what the SDK asked for but a narrowing of
// it — see cacheControl.
func PrivateCache(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		next.ServeHTTP(&cacheWriter{ResponseWriter: w}, r)
	})
}

// cacheWriter is the http.ResponseWriter [PrivateCache] hands down: it sets the
// two headers when the response is committed, whether that is an explicit
// WriteHeader or the implicit one of the first Write.
type cacheWriter struct {
	http.ResponseWriter
	committed bool
}

func (w *cacheWriter) WriteHeader(status int) {
	w.commit()
	w.ResponseWriter.WriteHeader(status)
}

func (w *cacheWriter) Write(b []byte) (int, error) {
	w.commit()
	return w.ResponseWriter.Write(b)
}

// Flush commits before flushing so a handler that streams without ever calling
// WriteHeader still leaves with the headers. http.NewResponseController prefers a
// Flush on the writer it is handed over one reached through Unwrap, so this
// method is what it finds — without it the flush would commit the response at the
// writer below and the two headers would never be written.
func (w *cacheWriter) Flush() {
	w.commit()
	//nolint:errcheck // http.Flusher.Flush reports nothing; the controller's error
	// is only about the writer not supporting flush, which Unwrap guarantees it does.
	_ = http.NewResponseController(w.ResponseWriter).Flush()
}

func (w *cacheWriter) commit() {
	if w.committed {
		return
	}
	w.committed = true
	w.Header().Set("Cache-Control", cacheControl)
	w.Header().Set("Vary", vary)
}

// Unwrap is load-bearing, not boilerplate: it is what keeps the streamable
// transport working through this wrapper. http.NewResponseController follows
// Unwrap to reach the real writer's Flush, and an SSE stream that could not be
// flushed would be a response no client sees until the handler returns — which,
// for a stream, is a response nobody sees at all. Deleting this method breaks
// streaming while leaving every status code and header test green, so
// TestUnwrapReachesTheUnderlyingFlusher exists to fail instead.
func (w *cacheWriter) Unwrap() http.ResponseWriter { return w.ResponseWriter }