package api

import (
	"errors"
	"net/http"
	"sync/atomic"
	"testing"
	"time"
)

// stubResponseWriter is the smallest possible http.ResponseWriter
// the leaf-detection helpers can be tested against. It deliberately
// does NOT implement Flusher so test scenarios can compose it under
// Flusher-faking wrappers to prove leafFlusherOf walks past them.
type stubResponseWriter struct {
	header http.Header
}

func newStubResponseWriter() *stubResponseWriter {
	return &stubResponseWriter{header: http.Header{}}
}

func (s *stubResponseWriter) Header() http.Header       { return s.header }
func (s *stubResponseWriter) Write([]byte) (int, error) { return 0, nil }
func (s *stubResponseWriter) WriteHeader(int)           {}

// flusherWrapper implements http.Flusher (so a single-level check
// would pass) AND Unwrap (so leafFlusherOf can walk past it). The
// flushes count proves whether the LEAF or the wrapper's no-op Flush
// got invoked.
type flusherWrapper struct {
	http.ResponseWriter
	flushes atomic.Int32
}

func (w *flusherWrapper) Flush() { w.flushes.Add(1) }
func (w *flusherWrapper) Unwrap() http.ResponseWriter {
	return w.ResponseWriter
}

// recordingFlusher implements Flusher directly so tests can prove
// leafFlusherOf returns the leaf itself when it is the leaf. The
// flushes pointer lets an integration test wrap a real
// http.ResponseWriter (forwarding writes) while still counting how
// many times the gateway reached the leaf via leafFlusherOf +
// leaf.Flush().
type recordingFlusher struct {
	http.ResponseWriter
	flushes    atomic.Int32
	flushesPtr *atomic.Int32
}

func (r *recordingFlusher) Flush() {
	r.flushes.Add(1)
	if r.flushesPtr != nil {
		r.flushesPtr.Add(1)
	}
}

// nilUnwrapperWrapper is a wrapper whose Unwrap returns nil. Tests
// that leafFlusherOf treats nil as terminal and falls back to the
// current writer as leaf.
type nilUnwrapperWrapper struct {
	http.ResponseWriter
}

func (w *nilUnwrapperWrapper) Flush()                      {}
func (w *nilUnwrapperWrapper) Unwrap() http.ResponseWriter { return nil }

// noUnwrapFlusher is a Fluzzer-implementation-only writer with NO
// Unwrap method. leafFlusherOf must treat it as a leaf.
type noUnwrapFlusher struct {
	http.ResponseWriter
	flushes atomic.Int32
}

func (n *noUnwrapFlusher) Flush() { n.flushes.Add(1) }

// deadlineCapableWriter is the minimum writer that lets
// http.NewResponseController succeed for SetReadDeadline. The
// implementation just records the most recent deadline.
type deadlineCapableWriter struct {
	http.ResponseWriter
	readDeadline atomic.Pointer[time.Time]
}

func (w *deadlineCapableWriter) SetReadDeadline(dl time.Time) error {
	w.readDeadline.Store(&dl)
	return nil
}

// deadlineRefusingWriter returns a sentinel error from
// SetReadDeadline to model a transport (e.g. an in-memory
// httptest.ResponseRecorder) that cannot honor connection-level
// deadlines.
type deadlineRefusingWriter struct {
	http.ResponseWriter
}

func (w *deadlineRefusingWriter) SetReadDeadline(time.Time) error {
	return http.ErrNotSupported
}

// TestLeafFlusherOfNil verifies the helper handles nil safely.
// A nil writer means "no leaf exists", so the helper must return
// (nil, false) — never panic — so the SSE pre-commit path can
// degrade to a stable 5xx instead of crashing the handler.
func TestLeafFlusherOfNil(t *testing.T) {
	f, ok := leafFlusherOf(nil)
	if ok || f != nil {
		t.Fatalf("leafFlusherOf(nil) = (%v, %v); want (nil, false)", f, ok)
	}
}

// TestLeafFlusherOfDirectFlusher verifies the helper returns the
// writer itself when it implements Flusher and has no Unwrap.
// This is the common case for net/http's built-in response writer.
func TestLeafFlusherOfDirectFlusher(t *testing.T) {
	leaf := &noUnwrapFlusher{ResponseWriter: newStubResponseWriter()}
	f, ok := leafFlusherOf(leaf)
	if !ok {
		t.Fatalf("ok = false; want true for direct Flusher")
	}
	f.Flush()
	if leaf.flushes.Load() != 1 {
		t.Fatalf("leaf flushes = %d; want 1 (Flusher must hit the leaf)", leaf.flushes.Load())
	}
}

// TestLeafFlusherOfDirectNonFlusher verifies the helper returns
// (nil, false) when the writer is a leaf and does NOT implement
// Flusher.
func TestLeafFlusherOfDirectNonFlusher(t *testing.T) {
	leaf := newStubResponseWriter()
	f, ok := leafFlusherOf(leaf)
	if ok || f != nil {
		t.Fatalf("leafFlusherOf(non-Flusher leaf) = (%v, %v); want (nil, false)", f, ok)
	}
}

// TestLeafFlusherOfUnwrapsOneLevelToFlusher verifies the helper
// descends a single Unwrap to find the leaf Flusher. Critical
// because gin's ResponseWriter always implements Flusher — a
// one-level unwrap is the minimum useful check.
func TestLeafFlusherOfUnwrapsOneLevelToFlusher(t *testing.T) {
	leaf := &noUnwrapFlusher{ResponseWriter: newStubResponseWriter()}
	wrapped := &flusherWrapper{ResponseWriter: leaf}
	f, ok := leafFlusherOf(wrapped)
	if !ok {
		t.Fatalf("ok = false; want true")
	}
	f.Flush()
	if wrapped.flushes.Load() != 0 {
		t.Fatalf("wrapper flushes = %d; want 0 (Flusher must NOT hit the wrapper)", wrapped.flushes.Load())
	}
	if leaf.flushes.Load() != 1 {
		t.Fatalf("leaf flushes = %d; want 1 (Flusher must hit the leaf)", leaf.flushes.Load())
	}
}

// TestLeafFlusherOfTwoFakeWrappersHidesNonFlusherLeaf is the
// critical anti-misjudgment test from reviewer3/codex5: a two-layer
// wrapper where BOTH intermediates implement Flusher but the leaf
// does NOT. A single-level unwrap against the outer wrapper would
// falsely conclude "Flusher available" and the gateway would
// commit a 200 + buffer SSE chunks forever. Recursive unwrap to
// the leaf is the only sound answer.
func TestLeafFlusherOfTwoFakeWrappersHidesNonFlusherLeaf(t *testing.T) {
	leaf := newStubResponseWriter()
	inner := &flusherWrapper{ResponseWriter: leaf}
	outer := &flusherWrapper{ResponseWriter: inner}
	f, ok := leafFlusherOf(outer)
	if ok || f != nil {
		t.Fatalf("leafFlusherOf(two-fake-wrapper → non-Flusher leaf) = (%v, %v); want (nil, false)", f, ok)
	}
}

// TestLeafFlusherOfTwoFakeWrappersReachesFlusherLeaf verifies the
// mirror case: when the leaf DOES implement Flusher, the helper
// unwraps past both fakes and returns the leaf's Flusher. This is
// the "real leaf Flusher preserves per-chunk gate" path.
func TestLeafFlusherOfTwoFakeWrappersReachesFlusherLeaf(t *testing.T) {
	leaf := &noUnwrapFlusher{ResponseWriter: newStubResponseWriter()}
	inner := &flusherWrapper{ResponseWriter: leaf}
	outer := &flusherWrapper{ResponseWriter: inner}
	f, ok := leafFlusherOf(outer)
	if !ok || f == nil {
		t.Fatalf("leafFlusherOf(two-fake-wrapper → Flusher leaf) = (%v, %v); want (leaf, true)", f, ok)
	}
	f.Flush()
	if outer.flushes.Load() != 0 || inner.flushes.Load() != 0 {
		t.Fatalf("intermediate flushes = outer=%d inner=%d; want 0/0", outer.flushes.Load(), inner.flushes.Load())
	}
	if leaf.flushes.Load() != 1 {
		t.Fatalf("leaf flushes = %d; want 1 (must hit the leaf, not an intermediate)", leaf.flushes.Load())
	}
}

// TestLeafFlusherOfUnwrapReturnsNilTerminates verifies the helper
// treats Unwrap()==nil as terminal and uses the current writer as
// the leaf. Some adapters (e.g. an http.Handler wrapping the
// writer to intercept Writes but not to expose a different writer)
// signal "no further wrapping" by returning nil from Unwrap.
func TestLeafFlusherOfUnwrapReturnsNilTerminates(t *testing.T) {
	// outer has Flush+Unwrap-returns-nil; the outer's Flush is what
	// should be returned.
	outer := &nilUnwrapperWrapper{ResponseWriter: newStubResponseWriter()}
	f, ok := leafFlusherOf(outer)
	if !ok {
		t.Fatalf("ok = false; want true (Unwrap()==nil means outer is the leaf)")
	}
	f.Flush()
}

// TestMaybeSetReadDeadlineCapableWriter verifies that when the
// underlying writer implements SetReadDeadline, the helper returns
// nil and the deadline actually lands.
func TestMaybeSetReadDeadlineCapableWriter(t *testing.T) {
	w := &deadlineCapableWriter{ResponseWriter: newStubResponseWriter()}
	dl := time.Now().Add(5 * time.Second)
	if err := maybeSetReadDeadline(w, dl); err != nil {
		t.Fatalf("maybeSetReadDeadline: %v", err)
	}
	if got := w.readDeadline.Load(); got == nil || !got.Equal(dl) {
		t.Fatalf("recorded deadline = %v; want %v", got, dl)
	}
}

type nonComparableWriter struct {
	http.ResponseWriter
	marker []byte
}

func (w nonComparableWriter) Flush() {}
func (w nonComparableWriter) Unwrap() http.ResponseWriter {
	return w
}

func TestLeafFlusherOfNonComparableCycleFailsClosedWithoutPanic(t *testing.T) {
	writer := nonComparableWriter{
		ResponseWriter: newStubResponseWriter(),
		marker:         []byte("non-comparable"),
	}
	var (
		flusher http.Flusher
		ok      bool
	)
	func() {
		defer func() {
			if recovered := recover(); recovered != nil {
				t.Fatalf("leafFlusherOf panicked for non-comparable writer: %v", recovered)
			}
		}()
		flusher, ok = leafFlusherOf(writer)
	}()
	if ok || flusher != nil {
		t.Fatalf("leafFlusherOf(non-comparable cycle) = (%v, %v), want (nil, false)", flusher, ok)
	}
}

// TestMaybeSetReadDeadlineNonCapableWriterReturnsErrNotSupported
// verifies that when the underlying writer does NOT support
// SetReadDeadline, the helper returns http.ErrNotSupported. The
// gateway MUST treat this as "cannot force-interrupt the in-flight
// read" and fall back to bounded shutdown.
func TestMaybeSetReadDeadlineNonCapableWriterReturnsErrNotSupported(t *testing.T) {
	w := &deadlineRefusingWriter{ResponseWriter: newStubResponseWriter()}
	err := maybeSetReadDeadline(w, time.Now())
	if !errors.Is(err, http.ErrNotSupported) {
		t.Fatalf("err = %v; want http.ErrNotSupported", err)
	}
}

// cyclicFlusherWrapper returns itself from Unwrap, creating an
// infinite A→A→A… cycle. leafFlusherOf's cycle guard must terminate
// the walk instead of spinning forever.
type cyclicFlusherWrapper struct {
	http.ResponseWriter
}

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

// cyclicNonFlusherWrapper is the same cycle but without Flush(),
// proving the guard also terminates when the leaf-candidate is a
// non-Flusher.
type cyclicNonFlusherWrapper struct {
	http.ResponseWriter
}

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

// TestLeafFlusherOfCyclicUnwrapTerminates verifies the cycle guard
// prevents an infinite agentteam when a wrapper's Unwrap returns itself.
// Without the guard, the for-agentteam would spin forever and the test
// would time out.
//
// Fail-closed: when a cycle is detected (or maxUnwrapDepth exceeded),
// the helper returns (nil, false). The cyclic writer's Flusher —
// even if it implements http.Flusher — is a fake; the gateway must
// NOT commit a 200 whose per-chunk Flush would land on a no-op stub.
// Both cases therefore want (nil, false), not (fake-Flusher, true).
func TestLeafFlusherOfCyclicUnwrapTerminates(t *testing.T) {
	cases := []struct {
		name string
		w    http.ResponseWriter
	}{
		{"cyclic-with-Flusher", &cyclicFlusherWrapper{ResponseWriter: newStubResponseWriter()}},
		{"cyclic-without-Flusher", &cyclicNonFlusherWrapper{ResponseWriter: newStubResponseWriter()}},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			done := make(chan struct{})
			var f http.Flusher
			var ok bool
			go func() {
				defer close(done)
				f, ok = leafFlusherOf(tc.w)
			}()
			select {
			case <-done:
			case <-time.After(2 * time.Second):
				t.Fatal("leafFlusherOf hung on cyclic Unwrap chain (cycle guard missing)")
			}
			if ok {
				t.Fatalf("ok = true; want false (cycle must fail-closed to (nil,false) so the gateway rejects pre-commit)")
			}
			if f != nil {
				t.Fatalf("f = %v; want nil (cycle must not return the fake Flusher)", f)
			}
		})
	}
}

// TestLeafFlusherOfMaxDepthFailClosed verifies the depth cap also
// fails closed instead of returning whatever writer the agentteam landed
// on. A chain longer than maxUnwrapDepth is almost certainly a buggy
// or hostile adapter; treating the tail as a leaf would let it
// masquerade as a Flusher.
func TestLeafFlusherOfMaxDepthFailClosed(t *testing.T) {
	// Build a chain of depth maxUnwrapDepth+1 by chaining
	// flusherWrappers. Each link Unwrap()s to the next. The leaf is
	// a real Flusher, but the helper must reject because the chain
	// exceeds the depth cap.
	const overshoot = 5
	leaf := &noUnwrapFlusher{ResponseWriter: newStubResponseWriter()}
	var wrapped http.ResponseWriter = leaf
	for i := 0; i < 64+overshoot; i++ {
		wrapped = &flusherWrapper{ResponseWriter: wrapped}
	}
	f, ok := leafFlusherOf(wrapped)
	if ok {
		t.Fatalf("ok = true; want false (chain exceeds maxUnwrapDepth, must fail-closed)")
	}
	if f != nil {
		t.Fatalf("f = %v; want nil", f)
	}
}

// Compile-time assertions to lock the test stub shapes.
var (
	_ http.ResponseWriter                       = (*stubResponseWriter)(nil)
	_ http.ResponseWriter                       = (*flusherWrapper)(nil)
	_ http.Flusher                              = (*flusherWrapper)(nil)
	_ interface{ Unwrap() http.ResponseWriter } = (*flusherWrapper)(nil)
	_ http.ResponseWriter                       = (*recordingFlusher)(nil)
	_ http.Flusher                              = (*recordingFlusher)(nil)
	_ http.ResponseWriter                       = (*nilUnwrapperWrapper)(nil)
	_ http.Flusher                              = (*nilUnwrapperWrapper)(nil)
	_ interface{ Unwrap() http.ResponseWriter } = (*nilUnwrapperWrapper)(nil)
	_ http.ResponseWriter                       = (*noUnwrapFlusher)(nil)
	_ http.Flusher                              = (*noUnwrapFlusher)(nil)
	_ http.ResponseWriter                       = (*deadlineCapableWriter)(nil)
	_ http.ResponseWriter                       = (*deadlineRefusingWriter)(nil)
	_ http.ResponseWriter                       = (*cyclicFlusherWrapper)(nil)
	_ http.Flusher                              = (*cyclicFlusherWrapper)(nil)
	_ interface{ Unwrap() http.ResponseWriter } = (*cyclicFlusherWrapper)(nil)
	_ http.ResponseWriter                       = (*cyclicNonFlusherWrapper)(nil)
	_ interface{ Unwrap() http.ResponseWriter } = (*cyclicNonFlusherWrapper)(nil)
	_ http.ResponseWriter                       = nonComparableWriter{}
	_ http.Flusher                              = nonComparableWriter{}
	_ interface{ Unwrap() http.ResponseWriter } = nonComparableWriter{}
)