package llm

// Observes how much of a gateway response the client will read. The body is read before
// the status is looked at, so a misbehaving gateway or proxy that streams forever was
// bounded only by the client's overall timeout (#53).

import (
	"context"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"
	"time"
)

// endlessGateway answers with the given status and streams filler until the client hangs
// up; wrote reports how many bytes it managed to send.
func endlessGateway(t *testing.T, status int, wrote *atomic.Int64) *httptest.Server {
	t.Helper()
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(status)
		chunk := []byte(strings.Repeat("a", 4096))
		for {
			n, err := w.Write(chunk)
			wrote.Add(int64(n))
			if err != nil {
				return
			}
			if f, ok := w.(http.Flusher); ok {
				f.Flush()
			}
		}
	}))
	t.Cleanup(srv.Close)
	return srv
}

func TestClient_ASuccessBodyLargerThanTheCapIsAnErrorAndTheClientStopsReading(t *testing.T) {
	var wrote atomic.Int64
	srv := endlessGateway(t, http.StatusOK, &wrote)
	c := New(srv.URL, "", "m", 30*time.Second, WithMaxResponseBytes(8192))

	start := time.Now()
	_, err := c.Complete(context.Background(), "s", "u")

	if err == nil || !strings.Contains(err.Error(), "exceeds") {
		t.Fatalf("want a response-too-large error, got %v", err)
	}
	if time.Since(start) > 5*time.Second {
		t.Fatalf("the client kept reading: took %s", time.Since(start))
	}
	// Hanging up lets the server's writes fail; allow what the transport had already buffered.
	time.Sleep(100 * time.Millisecond)
	if got := wrote.Load(); got > 8192+1<<20 {
		t.Fatalf("server managed to send %d bytes after a cap of 8192", got)
	}
}

func TestClient_AFailingBodyIsReadOnlyUpToTheErrorBound(t *testing.T) {
	var wrote atomic.Int64
	srv := endlessGateway(t, http.StatusBadGateway, &wrote)
	c := New(srv.URL, "", "m", 30*time.Second)

	start := time.Now()
	_, err := c.Complete(context.Background(), "s", "u")

	he, ok := err.(*HTTPError)
	if !ok {
		t.Fatalf("want *HTTPError, got %T: %v", err, err)
	}
	if he.Code != http.StatusBadGateway {
		t.Fatalf("code %d", he.Code)
	}
	if len(he.Body) > maxErrorBody+1 {
		t.Fatalf("kept %d bytes of an error body, bound is %d", len(he.Body), maxErrorBody)
	}
	if time.Since(start) > 5*time.Second {
		t.Fatalf("the client kept reading an endless error body: took %s", time.Since(start))
	}
}