package e2e

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"net/http"
	"os"
	"slices"
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

func bytesReader(b []byte) *bytes.Reader { return bytes.NewReader(b) }

func drainBody(resp *http.Response) {
	_, _ = io.Copy(io.Discard, resp.Body)
	_ = resp.Body.Close()
}

// TestConcurrentStreams is the first half of the phase 1 exit criterion:
// 1000 concurrent streams behind 3 sim-engines, all of them completing.
func TestConcurrentStreams(t *testing.T) {
	concurrency := envInt("LLMGW_E2E_CONCURRENCY", 1000)
	if testing.Short() {
		concurrency = 50
	}
	const tokens = 16

	h := start(t, 3)

	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
	defer cancel()

	var (
		ok       atomic.Int64
		short    atomic.Int64
		failed   atomic.Int64
		firstErr atomic.Value
		wg       sync.WaitGroup
	)

	start := time.Now()
	for i := range concurrency {
		wg.Add(1)
		go func() {
			defer wg.Done()
			n, _, err := streamRequest(ctx, h.router.URL,
				chatBody(fmt.Sprintf("request %d", i), tokens, true),
				map[string]string{"x-session-id": fmt.Sprintf("s%d", i%97)})
			switch {
			case err != nil:
				failed.Add(1)
				firstErr.CompareAndSwap(nil, err)
			case n < tokens:
				short.Add(1)
			default:
				ok.Add(1)
			}
		}()
	}
	wg.Wait()
	elapsed := time.Since(start)

	t.Logf("%d streams in %v (%.0f req/s): ok=%d short=%d failed=%d",
		concurrency, elapsed.Round(time.Millisecond),
		float64(concurrency)/elapsed.Seconds(), ok.Load(), short.Load(), failed.Load())

	if f := failed.Load(); f > 0 {
		t.Fatalf("%d of %d streams failed; first error: %v", f, concurrency, firstErr.Load())
	}
	if s := short.Load(); s > 0 {
		t.Errorf("%d streams ended before %d tokens", s, tokens)
	}

	// Load must actually have been spread; a single engine taking everything
	// would pass the checks above while defeating the point of the router.
	used := 0
	for _, e := range h.engines {
		if st := e.eng.Stats(); st.PrefixQueries > 0 || st.CachedBlocks > 0 {
			used++
		}
	}
	if used < 2 {
		t.Errorf("only %d of %d engines saw traffic", used, len(h.engines))
	}
}

// TestClientDisconnectStopsUpstream is the second half of the exit criterion:
// when a client goes away, the GPU must stop working on its request.
func TestClientDisconnectStopsUpstream(t *testing.T) {
	const streams = 40

	h := start(t, 3)

	ctx, cancel := context.WithCancel(context.Background())
	var wg sync.WaitGroup
	for i := range streams {
		wg.Add(1)
		go func() {
			defer wg.Done()
			// Long generations, so every request is still running when cancelled.
			_, _, _ = streamRequest(ctx, h.router.URL,
				chatBody(fmt.Sprintf("long request %d", i), 100000, true), nil)
		}()
	}

	// Wait until the engines are actually busy before pulling the plug.
	deadline := time.Now().Add(10 * time.Second)
	for h.totalRunning() < streams && time.Now().Before(deadline) {
		time.Sleep(10 * time.Millisecond)
	}
	if busy := h.totalRunning(); busy < streams {
		t.Fatalf("only %d of %d requests reached the engines", busy, streams)
	}

	cancel()
	wg.Wait()

	// The engines must drain: cancellation has to propagate router -> engine.
	deadline = time.Now().Add(15 * time.Second)
	for time.Now().Before(deadline) {
		if h.totalRunning() == 0 {
			// And the KV blocks must come back, not just the request slots.
			for _, e := range h.engines {
				if u := e.eng.Stats().KVCacheUsage; u != 0 {
					t.Fatalf("KV cache still held after cancellation: %v", u)
				}
			}
			return
		}
		time.Sleep(20 * time.Millisecond)
	}
	t.Fatalf("engines still working %d requests after every client disconnected", h.totalRunning())
}

// TestRouterOverhead is the third part of the exit criterion: the router should
// add little latency of its own.
//
// It measures the same request against an engine directly and through the
// router, and reports the p99 difference. The assertion is deliberately loose,
// because a shared CI machine cannot hold a 2 ms bound reliably; set
// LLMGW_E2E_STRICT=1 on quiet hardware to enforce the plan's actual target.
func TestRouterOverhead(t *testing.T) {
	samples := envInt("LLMGW_E2E_SAMPLES", 300)
	if testing.Short() {
		samples = 50
	}

	h := start(t, 1)
	direct := h.engines[0].srv.URL
	body := chatBody("overhead probe", 1, false)

	measure := func(url string) []time.Duration {
		out := make([]time.Duration, 0, samples)
		for range samples {
			start := time.Now()
			resp, err := http.Post(url+"/v1/chat/completions", "application/json", bytesReader(body)) //nolint:noctx // test
			if err != nil {
				t.Fatal(err)
			}
			drainBody(resp)
			out = append(out, time.Since(start))
		}
		slices.Sort(out)
		return out
	}

	// Warm both paths so connection setup does not land in the samples.
	measure(direct)
	measure(h.router.URL)

	d := measure(direct)
	r := measure(h.router.URL)

	p := func(s []time.Duration, q float64) time.Duration {
		return s[min(int(float64(len(s))*q), len(s)-1)]
	}
	overheadP50 := p(r, 0.50) - p(d, 0.50)
	overheadP99 := p(r, 0.99) - p(d, 0.99)

	t.Logf("direct p50=%v p99=%v | routed p50=%v p99=%v | overhead p50=%v p99=%v",
		p(d, 0.50).Round(time.Microsecond), p(d, 0.99).Round(time.Microsecond),
		p(r, 0.50).Round(time.Microsecond), p(r, 0.99).Round(time.Microsecond),
		overheadP50.Round(time.Microsecond), overheadP99.Round(time.Microsecond))

	limit := 25 * time.Millisecond
	if os.Getenv("LLMGW_E2E_STRICT") == "1" {
		limit = 2 * time.Millisecond
	}
	if overheadP99 > limit {
		t.Errorf("router added %v at p99, over the %v budget", overheadP99, limit)
	}
}