package mcpsrv

import (
	"bufio"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// The wrapper privateCache hands down, tested where it can be reached. The
// endpoint-level statement — that every answer of /mcp carries these two headers
// — is in cache_test.go; what is here is the three ways a handler can commit a
// response, because a wrapper that catches only one of them passes every header
// test while setting no header at all on the path the SDK actually takes.

// TestHeadersLandOnAnExplicitWriteHeader is the ordinary path.
func TestHeadersLandOnAnExplicitWriteHeader(t *testing.T) {
	h := privateCache(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusAccepted)
	}))

	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, httptest.NewRequest(http.MethodPost, "/mcp", nil))

	assert.Equal(t, http.StatusAccepted, rec.Code)
	assert.Equal(t, cacheControl, rec.Header().Get("Cache-Control"))
	assert.Equal(t, cacheVary, rec.Header().Get("Vary"))
}

// TestHeadersLandOnAnImplicitCommit covers the handler that never calls
// WriteHeader at all. net/http commits on the first Write and drops every header
// set after that point, silently.
func TestHeadersLandOnAnImplicitCommit(t *testing.T) {
	h := privateCache(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		_, _ = w.Write([]byte(`{"jsonrpc":"2.0"}`))
	}))

	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, httptest.NewRequest(http.MethodPost, "/mcp", nil))

	assert.Equal(t, http.StatusOK, rec.Code)
	assert.Equal(t, cacheControl, rec.Header().Get("Cache-Control"))
	assert.Equal(t, cacheVary, rec.Header().Get("Vary"))
	assert.Equal(t, `{"jsonrpc":"2.0"}`, rec.Body.String())
}

// TestHeadersLandOnAFlush is the third commit, and the one a stream takes: an
// SSE handler writes nothing before its first flush, so a wrapper that hooks
// Write and WriteHeader alone leaves the response uncached-marked exactly on the
// answers that stay open longest.
func TestHeadersLandOnAFlush(t *testing.T) {
	h := privateCache(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		require.NoError(t, http.NewResponseController(w).Flush())
	}))

	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/mcp", nil))

	assert.True(t, rec.Flushed)
	assert.Equal(t, cacheControl, rec.Header().Get("Cache-Control"))
	assert.Equal(t, cacheVary, rec.Header().Get("Vary"))
}

// TestTheSDKsOwnDirectivesAreOverridden is why this is a wrapper and not a
// middleware that sets the headers on the way in. The streamable transport sets
// Cache-Control itself, with Set, from inside the handler, so anything written
// before it runs loses. The handler here does exactly what the SDK does.
func TestTheSDKsOwnDirectivesAreOverridden(t *testing.T) {
	h := privateCache(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.Header().Set("Cache-Control", "no-cache, no-transform")
		w.WriteHeader(http.StatusOK)
	}))

	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, httptest.NewRequest(http.MethodPost, "/mcp", nil))

	assert.Equal(t, cacheControl, rec.Header().Get("Cache-Control"),
		"no-cache lets a cache store the body and revalidate, which is the thing no-store forbids")
	assert.Equal(t, cacheVary, rec.Header().Get("Vary"))
}

// TestOnlyTheFirstCommitWrites checks the committed flag: a handler that writes
// after committing must not have its headers re-Set on every Write.
func TestOnlyTheFirstCommitWrites(t *testing.T) {
	h := privateCache(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write([]byte("a"))
		_, _ = w.Write([]byte("b"))
	}))

	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, httptest.NewRequest(http.MethodPost, "/mcp", nil))

	assert.Equal(t, []string{cacheControl}, rec.Header().Values("Cache-Control"))
	assert.Equal(t, []string{cacheVary}, rec.Header().Values("Vary"))
	assert.Equal(t, "ab", rec.Body.String())
}

// TestUnwrapReachesTheWriterBelow is the test the Unwrap method exists for, and
// it deliberately does not measure flushing.
//
// The sentence usually attached to Unwrap on a writer like this — that without
// it the flusher is hidden and SSE breaks — is not true of this type: Flush is a
// method on cacheWriter, so a controller finds that one first and never unwraps.
// (Measured, not assumed: with Unwrap deleted every other test in this file, and
// every MCP session the rest of this package opens, stays green.) What stops
// working is everything cacheWriter does not implement itself, and the deadlines
// are the reachable half of that — so this is written against a deadline, which
// is the call that actually goes ErrNotSupported.
//
// It runs over a real server because httptest.ResponseRecorder supports no
// deadline at all, and a recorder therefore cannot tell the two cases apart.
func TestUnwrapReachesTheWriterBelow(t *testing.T) {
	var deadlineErr error
	srv := httptest.NewServer(privateCache(http.HandlerFunc(
		func(w http.ResponseWriter, _ *http.Request) {
			deadlineErr = http.NewResponseController(w).SetWriteDeadline(time.Now().Add(time.Minute))
		})))
	t.Cleanup(srv.Close)

	req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, srv.URL, nil)
	require.NoError(t, err)
	resp, err := (&http.Client{Timeout: 5 * time.Second}).Do(req)
	require.NoError(t, err)
	t.Cleanup(func() { _ = resp.Body.Close() })

	require.NoError(t, deadlineErr, "a controller must reach the writer this one wraps")
}

// TestAStreamReachesTheClientBeforeTheHandlerReturns is the same property
// measured rather than asserted: over a real connection, with a real client,
// bytes flushed inside the handler have to arrive while the handler is still
// running. A recorder cannot tell a flush that worked from one that was buffered
// until the end — and a broken flush here is what would turn every MCP session
// on this endpoint into a hang rather than an error.
func TestAStreamReachesTheClientBeforeTheHandlerReturns(t *testing.T) {
	release := make(chan struct{})
	srv := httptest.NewServer(privateCache(http.HandlerFunc(
		func(w http.ResponseWriter, _ *http.Request) {
			w.Header().Set("Content-Type", "text/event-stream")
			_, _ = w.Write([]byte("data: first\n\n"))
			if err := http.NewResponseController(w).Flush(); err != nil {
				t.Errorf("flush: %v", err)
			}
			<-release
		})))
	t.Cleanup(func() {
		close(release)
		srv.Close()
	})

	req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, srv.URL, nil)
	require.NoError(t, err)
	resp, err := (&http.Client{Timeout: 5 * time.Second}).Do(req)
	require.NoError(t, err)
	t.Cleanup(func() { _ = resp.Body.Close() })

	assert.Equal(t, cacheControl, resp.Header.Get("Cache-Control"))
	assert.Equal(t, cacheVary, resp.Header.Get("Vary"))

	line, err := bufio.NewReader(resp.Body).ReadString('\n')
	require.NoError(t, err, "the first event must arrive while the handler is still blocked")
	assert.Equal(t, "data: first\n", line)
}