package middlewaredemo

import (
	"bytes"
	"compress/gzip"
	"context"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"sync"
	"testing"
	"time"
)

type measuredWriter struct {
	http.ResponseWriter
	status, bytes int
}

func (w *measuredWriter) Unwrap() http.ResponseWriter { return w.ResponseWriter }
func (w *measuredWriter) WriteHeader(code int) {
	if code >= 100 && code < 200 && code != http.StatusSwitchingProtocols {
		w.ResponseWriter.WriteHeader(code)
		return
	}
	if w.status != 0 {
		return
	}
	w.status = code
	w.ResponseWriter.WriteHeader(code)
}
func (w *measuredWriter) Write(p []byte) (int, error) {
	if w.status == 0 {
		w.WriteHeader(200)
	}
	n, e := w.ResponseWriter.Write(p)
	w.bytes += n
	return n, e
}
func (w *measuredWriter) FlushError() error {
	if w.status == 0 {
		w.WriteHeader(200)
	}
	return http.NewResponseController(w.ResponseWriter).Flush()
}
func accessLog(logger *slog.Logger, next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		mw := &measuredWriter{ResponseWriter: w}
		defer func() {
			status := mw.status
			if status == 0 {
				status = 200
			}
			logger.InfoContext(r.Context(), "request", "method", r.Method, "status", status, "bytes", mw.bytes)
		}()
		next.ServeHTTP(mw, r)
	})
}
func TestLoggingAndFlush(t *testing.T) {
	var out bytes.Buffer
	logger := slog.New(slog.NewJSONHandler(&out, &slog.HandlerOptions{ReplaceAttr: func(groups []string, a slog.Attr) slog.Attr {
		if a.Key == slog.TimeKey {
			return slog.Attr{}
		}
		return a
	}}))
	h := accessLog(logger, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		io.WriteString(w, "hello")
		if err := http.NewResponseController(w).Flush(); err != nil {
			t.Error(err)
		}
	}))
	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, httptest.NewRequest("GET", "/", nil))
	if !rec.Flushed {
		t.Fatal("not flushed")
	}
	t.Log(out.String())
	var plain http.ResponseWriter = &struct{ http.ResponseWriter }{rec}
	if err := http.NewResponseController(plain).Flush(); !errors.Is(err, http.ErrNotSupported) {
		t.Fatal(err)
	}
	t.Log("plain wrapper loses Flush; explicit FlushError keeps it")
}

type bucket struct {
	mu                          sync.Mutex
	tokens, capacity, perSecond float64
	last                        time.Time
}

func (b *bucket) allow(now time.Time) bool {
	b.mu.Lock()
	defer b.mu.Unlock()
	elapsed := now.Sub(b.last).Seconds()
	if elapsed > 0 {
		b.tokens = min(b.capacity, b.tokens+elapsed*b.perSecond)
		b.last = now
	}
	if b.tokens < 1 {
		return false
	}
	b.tokens--
	return true
}
func Example_bucket() {
	now := time.Unix(0, 0)
	b := bucket{tokens: 2, capacity: 2, perSecond: 1, last: now}
	fmt.Println(b.allow(now), b.allow(now), b.allow(now))
	fmt.Println(b.allow(now.Add(time.Second)))
	// Output:
	// true true false
	// true
}
func limited(b *bucket, next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if !b.allow(time.Now()) {
			http.Error(w, "rate limited", 429)
			return
		}
		next.ServeHTTP(w, r)
	})
}
func TestLimit(t *testing.T) {
	b := &bucket{tokens: 1, capacity: 1, perSecond: 0, last: time.Now()}
	h := limited(b, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(204) }))
	for _, want := range []int{204, 429} {
		w := httptest.NewRecorder()
		h.ServeHTTP(w, httptest.NewRequest("GET", "/", nil))
		if w.Code != want {
			t.Fatal(w.Code)
		}
	}
}

func TestContextAndResponseTimeout(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	r := httptest.NewRequest("GET", "/", nil).WithContext(ctx)
	w := httptest.NewRecorder()
	http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { <-r.Context().Done(); io.WriteString(w, "still ran") }).ServeHTTP(w, r)
	t.Logf("canceled context without response policy: %d %q", w.Code, w.Body.String())
	release := make(chan struct{})
	written := make(chan error, 1)
	h := http.TimeoutHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		<-release
		_, err := io.WriteString(w, "too late")
		written <- err
	}), 10*time.Millisecond, "timed out")
	w = httptest.NewRecorder()
	h.ServeHTTP(w, httptest.NewRequest("GET", "/", nil))
	if w.Code != 503 {
		t.Fatal(w.Code)
	}
	close(release)
	err := <-written
	if !errors.Is(err, http.ErrHandlerTimeout) {
		t.Fatal(err)
	}
	t.Logf("TimeoutHandler: %d; late write rejected=%v", w.Code, errors.Is(err, http.ErrHandlerTimeout))
}

type gzipWriter struct {
	http.ResponseWriter
	z *gzip.Writer
}

func (w *gzipWriter) Unwrap() http.ResponseWriter { return w.ResponseWriter }
func (w *gzipWriter) Write(p []byte) (int, error) { return w.z.Write(p) }
func (w *gzipWriter) FlushError() error {
	if err := w.z.Flush(); err != nil {
		return err
	}
	return http.NewResponseController(w.ResponseWriter).Flush()
}
func TestGzipFlushAndClose(t *testing.T) {
	rec := httptest.NewRecorder()
	rec.Header().Set("Content-Encoding", "gzip")
	rec.Header().Add("Vary", "Accept-Encoding")
	z := gzip.NewWriter(rec)
	w := &gzipWriter{ResponseWriter: rec, z: z}
	if _, err := io.WriteString(w, "hello"); err != nil {
		t.Fatal(err)
	}
	if err := http.NewResponseController(w).Flush(); err != nil {
		t.Fatal(err)
	}
	if !rec.Flushed {
		t.Fatal("not flushed")
	}
	if err := z.Close(); err != nil {
		t.Fatal(err)
	}
	r, err := gzip.NewReader(bytes.NewReader(rec.Body.Bytes()))
	if err != nil {
		t.Fatal(err)
	}
	defer r.Close()
	body, err := io.ReadAll(r)
	if err != nil || string(body) != "hello" {
		t.Fatalf("%q %v", body, err)
	}
	t.Log("gzip inner flush, HTTP flush, and final trailer verified")
}

// Header sink deliberately does not model network framing; these assertions
// test the wrappers' state transitions, including 101's terminal status.
type headerSink struct {
	h     http.Header
	codes []int
}

func (w *headerSink) Header() http.Header {
	if w.h == nil {
		w.h = make(http.Header)
	}
	return w.h
}
func (w *headerSink) WriteHeader(code int)        { w.codes = append(w.codes, code) }
func (w *headerSink) Write(p []byte) (int, error) { return len(p), nil }
func TestInformationalAndUpgradeState(t *testing.T) {
	mw := &measuredWriter{ResponseWriter: new(headerSink)}
	mw.WriteHeader(103)
	if mw.status != 0 {
		t.Fatal(mw.status)
	}
	mw.WriteHeader(200)
	if mw.status != 200 {
		t.Fatal(mw.status)
	}
	mw = &measuredWriter{ResponseWriter: new(headerSink)}
	mw.WriteHeader(101)
	mw.WriteHeader(500)
	if mw.status != 101 {
		t.Fatal(mw.status)
	}
	cw := &committedWriter{ResponseWriter: new(headerSink)}
	cw.WriteHeader(103)
	if cw.committed {
		t.Fatal("103 committed")
	}
	cw.WriteHeader(200)
	if !cw.committed {
		t.Fatal("200 not committed")
	}
	cw = &committedWriter{ResponseWriter: new(headerSink)}
	cw.WriteHeader(101)
	if !cw.committed {
		t.Fatal("101 not committed")
	}
	t.Log("103 then 200 tracked; 101 treated as terminal")
}