package server

import (
	"bufio"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"
	"time"
)

func TestClientIPTrust(t *testing.T) {
	s := &Server{}
	if err := s.SetTrustedProxies([]string{"10.0.0.0/8", "::1/128"}); err != nil {
		t.Fatal(err)
	}
	if err := s.SetTrustedProxies([]string{"bad"}); err == nil {
		t.Fatal("invalid CIDR accepted")
	}
	for _, tt := range []struct{ peer, xff, want string }{
		{"192.0.2.1:42", "198.51.100.1", "192.0.2.1"},
		{"10.0.0.1:42", "198.51.100.1, 10.0.0.2", "198.51.100.1"},
		{"10.0.0.1:42", "198.51.100.1, 192.0.2.2", "192.0.2.2"},
		{"10.0.0.1:42", "bad, 192.0.2.2", "10.0.0.1"},
		{"10.0.0.1:42", "", "10.0.0.1"},
		{"10.0.0.1:42", "10.0.0.2", "10.0.0.1"},
		{"[::1]:42", "2001:db8::1", "2001:db8::1"},
	} {
		r := httptest.NewRequest("POST", "/", nil)
		r.RemoteAddr = tt.peer
		r.Header.Set("X-Forwarded-For", tt.xff)
		r.Header.Set("X-Real-IP", "203.0.113.1")
		if got := s.clientIP(r); got != tt.want {
			t.Errorf("%+v: got %s", tt, got)
		}
	}
}

func TestLoginProxyIsolation(t *testing.T) {
	s, err := New(":memory:", "password", nil)
	if err != nil {
		t.Fatal(err)
	}
	defer s.Close()
	if err := s.SetTrustedProxies([]string{"10.0.0.0/8"}); err != nil {
		t.Fatal(err)
	}
	login := func(peer, forwarded string) int {
		r := httptest.NewRequest("POST", "/admin/api/login", strings.NewReader(`{"password":"wrong"}`))
		r.RemoteAddr = peer
		r.Header.Set("X-Forwarded-For", forwarded)
		w := httptest.NewRecorder()
		s.login(w, r)
		return w.Code
	}
	for i := 0; i < 10; i++ {
		if got := login("10.0.0.1:1", "192.0.2.1"); got != 401 {
			t.Fatal(got)
		}
	}
	if got := login("10.0.0.1:1", "192.0.2.1"); got != 429 {
		t.Fatal(got)
	}
	if got := login("10.0.0.1:1", "192.0.2.2"); got != 401 {
		t.Fatal(got)
	}
	for i := 0; i < 10; i++ {
		login("192.0.2.3:1", fmt.Sprintf("198.51.100.%d", i))
	}
	if got := login("192.0.2.3:1", "203.0.113.1"); got != 429 {
		t.Fatal(got)
	}
}

func TestBodyDeadline(t *testing.T) {
	result := make(chan error, 1)
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		err := readBody(w, r, 1024, 40*time.Millisecond, func() error { _, err := io.ReadAll(r.Body); return err })
		result <- err
		w.WriteHeader(400)
	}))
	defer srv.Close()
	conn, err := net.Dial("tcp", srv.Listener.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close()
	fmt.Fprint(conn, "POST / HTTP/1.1\r\nHost: test\r\nContent-Length: 20\r\n\r\nx")
	select {
	case err := <-result:
		if err == nil {
			t.Fatal("slow body accepted")
		}
	case <-time.After(2 * time.Second):
		t.Fatal("body deadline did not fire")
	}
}

func TestBodyDeadlineCleared(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		err := readBody(w, r, 1024, 40*time.Millisecond, func() error { _, err := io.ReadAll(r.Body); return err })
		if err != nil {
			http.Error(w, "read", 400)
			return
		}
		time.Sleep(100 * time.Millisecond)
		select {
		case <-r.Context().Done():
			http.Error(w, "canceled", 500)
			return
		default:
		}
		fmt.Fprint(w, "data: ok\n\n")
	}))
	defer srv.Close()
	conn, err := net.Dial("tcp", srv.Listener.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close()
	conn.SetDeadline(time.Now().Add(2 * time.Second))
	fmt.Fprint(conn, "POST / HTTP/1.1\r\nHost: test\r\nContent-Length: 1\r\n\r\nx")
	resp, err := http.ReadResponse(bufio.NewReader(conn), nil)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		t.Fatal(resp.StatusCode)
	}
}

func TestLoginBodyLimit(t *testing.T) {
	r := httptest.NewRequest("POST", "/admin/api/login", strings.NewReader(`{"password":"`+strings.Repeat("x", 16<<10)+`"}`))
	var v any
	if decode(httptest.NewRecorder(), r, &v) == nil {
		t.Fatal("oversize login accepted")
	}
}

func TestLimiterReleaseAndCleanup(t *testing.T) {
	now := time.Now()
	l := newRequestLimiter(2, 1, 2)
	a, _ := l.acquire("a", now)
	if a == nil {
		t.Fatal("denied")
	}
	if release, _ := l.acquire("a", now); release != nil {
		t.Fatal("per-key concurrency bypass")
	}
	b, _ := l.acquire("b", now)
	if b == nil {
		t.Fatal("denied")
	}
	if release, _ := l.acquire("c", now); release != nil {
		t.Fatal("global concurrency bypass")
	}
	a()
	a()
	b()
	a, _ = l.acquire("a", now)
	if a == nil {
		t.Fatal("not released")
	}
	a()
	if release, retry := l.acquire("a", now); release != nil || retry != "60" {
		t.Fatal("rpm bypass")
	}
	c, _ := l.acquire("c", now.Add(time.Minute))
	if c == nil {
		t.Fatal("window not reset")
	}
	c()
	if len(l.entries) != 1 {
		t.Fatalf("stale entries: %d", len(l.entries))
	}
	for i := 0; i < maxLimitEntries; i++ {
		l.entries[fmt.Sprint(i)] = &limitEntry{window: now.Add(time.Minute)}
	}
	if release, _ := l.acquire("overflow", now.Add(time.Minute)); release != nil {
		t.Fatal("map cap bypass")
	}
}

func TestLimiterConcurrent(t *testing.T) {
	l := newRequestLimiter(10000, 8, 64)
	var wg sync.WaitGroup
	for i := 0; i < 100; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			if release, _ := l.acquire("key", time.Now()); release != nil {
				release()
			}
		}()
	}
	wg.Wait()
	if l.active != 0 || l.entries["key"].active != 0 {
		t.Fatal("permit leak")
	}
}

func TestProxyLimitRejectsBeforeBodyAndLogging(t *testing.T) {
	s, err := New(":memory:", "test-password", nil)
	if err != nil {
		t.Fatal(err)
	}
	defer s.Close()
	if err := s.SetLimits(1, 1, 1); err != nil {
		t.Fatal(err)
	}
	cookie := loginCookie(t, s)
	ch := integrationChannel(t, s, cookie, "abc", "openai-chat", "http://127.0.0.1:1")
	key := integrationCreateKey(t, s, cookie, "limited", ch)
	request := func(body io.Reader) *httptest.ResponseRecorder {
		r := httptest.NewRequest("POST", "/v1/chat/completions", body)
		r.Header.Set("Authorization", "Bearer "+key.raw)
		w := httptest.NewRecorder()
		s.ServeHTTP(w, r)
		return w
	}
	if w := request(strings.NewReader("{}")); w.Code != 400 {
		t.Fatal(w.Code)
	}
	if s.limiter.active != 0 {
		t.Fatal("permit leaked after invalid request")
	}
	body := &observedReader{}
	w := request(body)
	if w.Code != 429 || w.Header().Get("Retry-After") == "" || body.read {
		t.Fatalf("status=%d read=%t", w.Code, body.read)
	}
	if err := s.flushLogs(); err != nil {
		t.Fatal(err)
	}
	var count int
	if err := s.db.QueryRow("SELECT count(*) FROM logs").Scan(&count); err != nil {
		t.Fatal(err)
	}
	if count != 1 {
		t.Fatalf("rejected request logged: %d logs", count)
	}
}

type observedReader struct{ read bool }

func (r *observedReader) Read(p []byte) (int, error) { r.read = true; return 0, io.EOF }

func TestLoginFailuresOnlyAndGlobalCap(t *testing.T) {
	s, err := New(":memory:", "password", nil)
	if err != nil {
		t.Fatal(err)
	}
	defer s.Close()
	if err := s.SetTrustedProxies([]string{"10.0.0.0/8"}); err != nil {
		t.Fatal(err)
	}
	login := func(peer, forwarded, password string) int {
		r := httptest.NewRequest("POST", "/admin/api/login", strings.NewReader(`{"password":"`+password+`"}`))
		r.RemoteAddr = peer
		if forwarded != "" {
			r.Header.Set("X-Forwarded-For", forwarded)
		}
		w := httptest.NewRecorder()
		s.login(w, r)
		return w.Code
	}
	for i := 0; i < 15; i++ {
		if got := login("10.0.0.1:1", "192.0.2.1", "password"); got != 200 {
			t.Fatalf("correct login %d rejected: %d", i, got)
		}
	}
	for i := 0; i < loginFailuresPerMinute; i++ {
		if got := login("10.0.0.1:1", "192.0.2.2", "wrong"); got != 401 {
			t.Fatalf("failure %d: %d", i, got)
		}
	}
	if got := login("10.0.0.1:1", "192.0.2.2", "wrong"); got != 429 {
		t.Fatal(got)
	}
	// The per-IP lockout is uniform: it must not confirm candidate passwords.
	if got := login("10.0.0.1:1", "192.0.2.2", "password"); got != 429 {
		t.Fatal("lockout leaks password validity:", got)
	}
	// A successful login clears its IP bucket, so fat-fingers do not accumulate.
	if got := login("10.0.0.1:1", "192.0.2.3", "wrong"); got != 401 {
		t.Fatal(got)
	}
	if got := login("10.0.0.1:1", "192.0.2.3", "password"); got != 200 {
		t.Fatal(got)
	}
	for i := 0; i < loginFailuresPerMinute; i++ {
		if got := login("10.0.0.1:1", "192.0.2.3", "wrong"); got != 401 {
			t.Fatalf("post-clear failure %d: %d", i, got)
		}
	}
	if got := login("10.0.0.1:1", "192.0.2.3", "wrong"); got != 429 {
		t.Fatal(got)
	}
	for i := 0; i < loginGlobalFailuresPerMinute; i++ {
		login("10.0.0.1:1", fmt.Sprintf("198.18.%d.%d", i/256, i%256), "wrong")
	}
	if got := login("10.0.0.1:1", "203.0.113.1", "wrong"); got != 429 {
		t.Fatal(got)
	}
	// The global cap fails closed too: aggregate brute force locks every
	// login out without revealing candidate-password validity.
	if got := login("10.0.0.1:1", "203.0.113.1", "password"); got != 429 {
		t.Fatal("global lockout leaks password validity:", got)
	}
}

func TestAuthFailureGate(t *testing.T) {
	s := testServer(t)
	cookie := loginCookie(t, s)
	ch := integrationChannel(t, s, cookie, "abc", "openai-chat", "http://127.0.0.1:1")
	key := integrationCreateKey(t, s, cookie, "gated", ch)
	request := func(auth string) *httptest.ResponseRecorder {
		r := httptest.NewRequest("POST", "/v1/chat/completions", strings.NewReader(integrationRequest("openai-chat", "abc/gpt-5", false)))
		r.Header.Set("Authorization", "Bearer "+auth)
		w := httptest.NewRecorder()
		s.ServeHTTP(w, r)
		return w
	}
	if w := request(key.raw); w.Code != 502 {
		t.Fatal("valid key setup:", w.Code, w.Body.String())
	}
	for i := 0; i < authFailuresPerMinute; i++ {
		if w := request("wrong"); w.Code != 401 {
			t.Fatalf("%d: %d", i, w.Code)
		}
	}
	if w := request("wrong"); w.Code != 429 || w.Header().Get("Retry-After") == "" {
		t.Fatal(w.Code, w.Header())
	}
	// The gate applies before the key lookup: a valid key from the same
	// (attacker-shared) IP is throttled the same way.
	if w := request(key.raw); w.Code != 429 {
		t.Fatal(w.Code)
	}
}

func TestProxyBodyTooLarge(t *testing.T) {
	s := testServer(t)
	cookie := loginCookie(t, s)
	ch := integrationChannel(t, s, cookie, "abc", "openai-chat", "http://127.0.0.1:1")
	key := integrationCreateKey(t, s, cookie, "limit", ch)
	r := httptest.NewRequest("POST", "/v1/chat/completions", strings.NewReader(`{"model":"abc/gpt-5","messages":[],"pad":"`+strings.Repeat("x", maxBody)+`"}`))
	r.Header.Set("Authorization", "Bearer "+key.raw)
	w := httptest.NewRecorder()
	s.ServeHTTP(w, r)
	if w.Code != 413 || !strings.Contains(w.Body.String(), "too large") {
		t.Fatal(w.Code, w.Body.String())
	}
}