package proxy_test

import (
	"context"
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"git.wntrmute.dev/kyle/crossbar/internal/store"
)

// A client that goes away mid-stream is still a request that happened: it held a slot, it cost
// prefill, and it belongs in the accounting. The row records status 499 and a non-empty err.
func TestClientCancelMidStreamIsRecorded(t *testing.T) {
	slow := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/health":
			fmt.Fprint(w, `{"status":"ok"}`)
		case "/v1/models":
			fmt.Fprint(w, `{"object":"list","data":[{"id":"shared"}]}`)
		default:
			w.Header().Set("Content-Type", "text/event-stream")
			w.WriteHeader(200)
			fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"content\":\"first\"}}]}\n\n")
			w.(http.Flusher).Flush()
			select {
			case <-r.Context().Done():
			case <-time.After(3 * time.Second):
			}
		}
	}))
	t.Cleanup(slow.Close)
	beta := newUpstream(t, "beta")
	r := newRig(t, twoHosts, &upstream{name: "alpha", srv: slow}, beta)

	ctx, cancel := context.WithCancel(context.Background())
	body := `{"model":"alpha-only","stream":true,"messages":[{"role":"user","content":"cancel me"}]}`
	req, _ := http.NewRequestWithContext(ctx, http.MethodPost, r.front.URL+"/r/v1/chat/completions", strings.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	buf := make([]byte, 64)
	if _, err := resp.Body.Read(buf); err != nil {
		t.Fatalf("first chunk: %v", err)
	}
	cancel()
	resp.Body.Close()

	deadline := time.Now().Add(3 * time.Second)
	var counts []store.StatusCount
	for time.Now().Before(deadline) {
		counts, _ = r.store.StatusCounts(time.Time{})
		if len(counts) > 0 {
			break
		}
		time.Sleep(25 * time.Millisecond)
	}
	if len(counts) != 1 || counts[0].Status != 499 || counts[0].Route != "r" || counts[0].Count != 1 {
		t.Fatalf("status counts after a cancelled stream = %+v, want one row: route r, status 499", counts)
	}
	rows, _ := r.store.Usage(time.Time{}, store.ByRoute)
	if len(rows) != 1 || rows[0].Requests != 1 || rows[0].Errors != 1 {
		t.Errorf("usage = %+v, want 1 request counted as an error", rows)
	}
}

// The same when the client gives up while waiting in the queue: a 499 row, no slot leaked.
func TestClientCancelWhileQueuedIsRecorded(t *testing.T) {
	alpha := newUpstream(t, "alpha")
	alpha.delay = 800 * time.Millisecond
	r := newRig(t, `
listen = "127.0.0.1:1"
queue_max = 2
[hosts.alpha]
base_url = %q
models = { "shared" = { parallel = 1 } }
[routes.r]
hosts = ["alpha"]
default_model = "shared"
`, alpha)
	go func() { drain(r.post("/r/v1/chat/completions", conversation(1, 1))) }() // holds the one slot
	waitUntil(t, func() bool { return r.lim.InFlight("alpha", "shared") == 1 })
	ctx, cancel := context.WithTimeout(context.Background(), 150*time.Millisecond)
	defer cancel()
	req, _ := http.NewRequestWithContext(ctx, http.MethodPost, r.front.URL+"/r/v1/chat/completions", strings.NewReader(conversation(2, 1)))
	req.Header.Set("Content-Type", "application/json")
	if _, err := http.DefaultClient.Do(req); err == nil {
		t.Fatal("the queued request should have been cancelled by its context")
	}
	deadline := time.Now().Add(3 * time.Second)
	for time.Now().Before(deadline) {
		counts, _ := r.store.StatusCounts(time.Time{})
		for _, c := range counts {
			if c.Status == 499 {
				if r.lim.Queued("alpha", "shared") != 0 {
					t.Errorf("queued = %d after the waiter cancelled", r.lim.Queued("alpha", "shared"))
				}
				return
			}
		}
		time.Sleep(25 * time.Millisecond)
	}
	t.Fatal("no 499 row recorded for the request cancelled while queued")
}