package proxy

import (
	"context"
	"net/http"
	"net/http/httptest"
	"net/url"
	"testing"
)

func TestMihomoForwardRewritesPathAndAuth(t *testing.T) {
	target, err := url.Parse("http://127.0.0.1:9090")
	if err != nil {
		t.Fatal(err)
	}
	cap := &captureTransport{}
	h := NewMihomoForward(target, "/api/mtg/mihomo", "Bearer testsecret", cap, nil)

	req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, "http://gw/api/mtg/mihomo/proxies", nil)
	if err != nil {
		t.Fatal(err)
	}
	req.Header.Set("Authorization", "Bearer client-should-not-forward")
	req.Header.Set("Host", "public-gateway.example:8888")
	h.ServeHTTP(httptest.NewRecorder(), req)

	if cap.got == nil {
		t.Fatal("no outgoing request captured")
	}
	if got, want := cap.got.Host, "127.0.0.1:9090"; got != want {
		t.Fatalf("Host: got %q want %q (upstream must not see client Host)", got, want)
	}
	if got := cap.got.Header.Get("Authorization"); got != "Bearer testsecret" {
		t.Fatalf("Authorization: got %q want Bearer testsecret", got)
	}
	want, err := url.Parse("http://127.0.0.1:9090/proxies")
	if err != nil {
		t.Fatal(err)
	}
	assertSameURL(t, cap.got.URL, want)
}

func TestIsWebSocketUpgrade(t *testing.T) {
	t.Run("standard headers", func(t *testing.T) {
		r := httptest.NewRequest(http.MethodGet, "/ws", nil)
		r.Header.Set("Connection", "Upgrade")
		r.Header.Set("Upgrade", "websocket")
		if !isWebSocketUpgrade(r) {
			t.Fatal("expected websocket upgrade")
		}
	})

	t.Run("tokenized Connection", func(t *testing.T) {
		r := httptest.NewRequest(http.MethodGet, "/ws", nil)
		r.Header.Set("Connection", "keep-alive, Upgrade")
		r.Header.Set("Upgrade", "websocket")
		if !isWebSocketUpgrade(r) {
			t.Fatal("expected websocket upgrade for tokenized Connection")
		}
	})

	t.Run("Sec-WebSocket-Key fallback", func(t *testing.T) {
		r := httptest.NewRequest(http.MethodGet, "/ws", nil)
		r.Header.Set("Sec-WebSocket-Key", "dGhlIHNhbXBsZSBub25jZQ==")
		if !isWebSocketUpgrade(r) {
			t.Fatal("expected websocket upgrade with Sec-WebSocket-Key")
		}
	})

	t.Run("no upgrade headers", func(t *testing.T) {
		r := httptest.NewRequest(http.MethodGet, "/ws", nil)
		if isWebSocketUpgrade(r) {
			t.Fatal("plain GET should not be detected as websocket")
		}
	})

	t.Run("Upgrade without Connection", func(t *testing.T) {
		r := httptest.NewRequest(http.MethodGet, "/ws", nil)
		r.Header.Set("Upgrade", "websocket")
		if isWebSocketUpgrade(r) {
			t.Fatal("should not match without Connection header or Sec-WebSocket-Key")
		}
	})
}

func TestBuildWSUpgradeRequest(t *testing.T) {
	orig := httptest.NewRequest(http.MethodGet, "http://gw/api/mtg/mihomo/traffic", nil)
	orig.Header.Set("Connection", "Upgrade")
	orig.Header.Set("Upgrade", "websocket")
	orig.Header.Set("Sec-WebSocket-Version", "13")
	orig.Header.Set("Sec-WebSocket-Key", "dGhlIHNhbXBsZSBub25jZQ==")
	orig.Header.Set("Authorization", "Bearer client-token-must-not-leak")

	dest, _ := url.Parse("http://172.20.0.2:9090/traffic")
	raw := string(buildWSUpgradeRequest(orig, dest, "Bearer upstream-secret"))

	for _, want := range []string{
		"GET /traffic HTTP/1.1\r\n",
		"Host: 172.20.0.2:9090\r\n",
		"Connection: Upgrade\r\n",
		"Upgrade: websocket\r\n",
		"Sec-WebSocket-Version: 13\r\n",
		"Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n",
		"Authorization: Bearer upstream-secret\r\n",
		"\r\n",
	} {
		if !containsStr(raw, want) {
			t.Errorf("request missing %q\ngot:\n%s", want, raw)
		}
	}
	if containsStr(raw, "client-token-must-not-leak") {
		t.Error("client Authorization leaked into upstream request")
	}
}

func containsStr(s, sub string) bool {
	return len(s) >= len(sub) && (s == sub || len(sub) == 0 ||
		(len(s) > 0 && len(sub) > 0 && stringContains(s, sub)))
}

func stringContains(s, sub string) bool {
	for i := 0; i <= len(s)-len(sub); i++ {
		if s[i:i+len(sub)] == sub {
			return true
		}
	}
	return false
}