package server

import (
	"ai-router/internal/protocol"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"runtime"
	"sort"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

// Explicitly opt in: ROUTER_PERF=1 go test ./internal/server -run TestPerformanceHTTP -v -count=1.
func TestPerformanceHTTP(t *testing.T) {
	if os.Getenv("ROUTER_PERF") != "1" {
		t.Skip("opt-in load test")
	}
	for _, stream := range []bool{false, true} {
		up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			io.Copy(io.Discard, r.Body)
			r.Body.Close()
			if stream {
				w.Header().Set("Content-Type", "text/event-stream")
				for _, f := range strings.Split(integrationSSE("openai-chat"), "\n\n") {
					if f != "" {
						fmt.Fprint(w, f+"\n\n")
						w.(http.Flusher).Flush()
					}
				}
			} else {
				w.Header().Set("Content-Type", "application/json")
				fmt.Fprint(w, integrationResponse("openai-chat"))
			}
		}))
		for _, mode := range []string{"direct", "memory", "disk"} {
			for _, down := range []string{"openai-chat", "anthropic"} {
				if mode == "direct" && down != "openai-chat" {
					continue
				}
				for _, concurrency := range []int{1, 8, 32} {
					t.Run(fmt.Sprintf("stream=%t/%s/%s/c%d", stream, mode, down, concurrency), func(t *testing.T) {
						target := up.URL + integrationPath("openai-chat")
						raw := "test-key"
						var s *Server
						var router *httptest.Server
						if mode != "direct" {
							path := ":memory:"
							if mode == "disk" {
								path = filepath.Join(t.TempDir(), "router.db")
							}
							var err error
							s, err = New(path, "test-password", nil)
							if err != nil {
								t.Fatal(err)
							}
							defer s.Close()
							if err := s.SetLimits(1000000, 64, 128); err != nil {
								t.Fatal(err)
							}
							cookie := loginCookie(t, s)
							ch := integrationChannel(t, s, cookie, "abc", "openai-chat", up.URL)
							key := integrationCreateKey(t, s, cookie, "perf", ch)
							raw = key.raw
							router = httptest.NewServer(s)
							defer router.Close()
							target = router.URL + integrationPath(down)
						}
						tr := http.DefaultTransport.(*http.Transport).Clone()
						tr.MaxIdleConnsPerHost = 64
						defer tr.CloseIdleConnections()
						client := &http.Client{Transport: tr, Timeout: 10 * time.Second}
						body := integrationRequest(down, "abc/gpt-5", stream)
						request := func() (time.Duration, time.Duration, bool) {
							start := time.Now()
							r, _ := http.NewRequest("POST", target, strings.NewReader(body))
							r.Header.Set("Authorization", "Bearer "+raw)
							r.Header.Set("Content-Type", "application/json")
							resp, err := client.Do(r)
							if err != nil {
								return time.Since(start), 0, false
							}
							first := make([]byte, 1)
							_, err = resp.Body.Read(first)
							ttfb := time.Since(start)
							_, e := io.Copy(io.Discard, resp.Body)
							resp.Body.Close()
							return time.Since(start), ttfb, resp.StatusCode == 200 && (err == nil || err == io.EOF) && e == nil
						}
						for i := 0; i < 20; i++ {
							if _, _, ok := request(); !ok {
								t.Fatal("warmup failed")
							}
						}
						const n = 400
						dur := make([]float64, n)
						first := make([]float64, n)
						var index atomic.Int64
						var failures atomic.Int64
						var wg sync.WaitGroup
						runtime.GC()
						var before, after runtime.MemStats
						runtime.ReadMemStats(&before)
						start := time.Now()
						for c := 0; c < concurrency; c++ {
							wg.Add(1)
							go func() {
								defer wg.Done()
								for {
									i := int(index.Add(1)) - 1
									if i >= n {
										return
									}
									d, f, ok := request()
									dur[i] = float64(d) / 1e6
									first[i] = float64(f) / 1e6
									if !ok {
										failures.Add(1)
									}
								}
							}()
						}
						wg.Wait()
						elapsed := time.Since(start)
						runtime.ReadMemStats(&after)
						sort.Float64s(dur)
						sort.Float64s(first)
						var total float64
						for _, d := range dur {
							total += d
						}
						drained := elapsed
						if s != nil {
							router.Close()
							if err := s.flushLogs(); err != nil {
								t.Fatal(err)
							}
							drained = time.Since(start)
							var count int
							if err := s.db.QueryRow("SELECT COUNT(*) FROM logs").Scan(&count); err != nil || count != n+20 {
								t.Fatalf("log count %d err %v", count, err)
							}
						}
						t.Logf("RESULT rps=%.1f drained_rps=%.1f mean_ms=%.3f p50_ms=%.3f p95_ms=%.3f p99_ms=%.3f first_p50_ms=%.3f first_p95_ms=%.3f alloc_B_req=%d errors=%d", float64(n)/elapsed.Seconds(), float64(n)/drained.Seconds(), total/n, dur[n/2], dur[n*95/100], dur[n*99/100], first[n/2], first[n*95/100], (after.TotalAlloc-before.TotalAlloc)/n, failures.Load())
						if failures.Load() != 0 {
							t.Fail()
						}
					})
				}
			}
		}
		up.Close()
	}
}

func BenchmarkPerformanceConversion(b *testing.B) {
	for _, size := range []int{1024, 65536} {
		for _, to := range []protocol.Kind{protocol.Chat, protocol.Responses, protocol.Anthropic} {
			b.Run(fmt.Sprintf("%d/%s", size, to), func(b *testing.B) {
				body := []byte(`{"model":"abc/gpt-5","max_tokens":1024,"messages":[{"role":"user","content":"` + strings.Repeat("x", size) + `"}]}`)
				b.ReportAllocs()
				b.ResetTimer()
				for i := 0; i < b.N; i++ {
					if _, err := protocol.ConvertRequest(protocol.Chat, to, body, "gpt-4"); err != nil {
						b.Fatal(err)
					}
				}
			})
		}
	}
}

func BenchmarkPerformanceStats(b *testing.B) {
	for _, n := range []int{1000, 10000, 100000} {
		b.Run(fmt.Sprint(n), func(b *testing.B) {
			s, err := New(filepath.Join(b.TempDir(), "stats.db"), "test-password", nil)
			if err != nil {
				b.Fatal(err)
			}
			defer s.Close()
			_, err = s.db.Exec(`WITH RECURSIVE seq(x) AS (SELECT 1 UNION ALL SELECT x+1 FROM seq WHERE x<?) INSERT INTO logs(created_at,key_id,key_name,channel,request_model,actual_model,downstream,upstream,status,duration_ms,stream,input_tokens,output_tokens,cache_read_tokens,cache_write_tokens,usage_known,error,raw_usage) SELECT '2026-10-20T10:00:00Z','key','perf','abc','abc/gpt-5','gpt-4','openai-chat','openai-chat',200,10,0,1000,100,100,0,1,'','{}' FROM seq`, n)
			if err != nil {
				b.Fatal(err)
			}
			r := httptest.NewRequest("GET", "/admin/api/stats", nil)
			b.ReportAllocs()
			b.ResetTimer()
			for i := 0; i < b.N; i++ {
				w := httptest.NewRecorder()
				s.stats(w, r)
				if w.Code != 200 {
					b.Fatal(w.Code)
				}
			}
		})
	}
}