package lb

import (
	"bytes"
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"strconv"
	"strings"
	"sync"
	"testing"
	"time"

	coreapi "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/types"
)

// testProxy binds every requested address on loopback and remembers where.
type testProxy struct {
	*Proxy
	mu     sync.Mutex
	actual map[string]string
}

func newTestProxy(t *testing.T) *testProxy {
	tp := &testProxy{Proxy: NewProxy(), actual: map[string]string{}}
	tp.listen = func(addr string) (net.Listener, error) {
		ln, err := net.Listen("tcp", "127.0.0.1:0")
		if err == nil {
			tp.mu.Lock()
			tp.actual[addr] = ln.Addr().String()
			tp.mu.Unlock()
		}
		return ln, err
	}
	t.Cleanup(tp.Close)
	return tp
}

func (tp *testProxy) at(port int32) string {
	tp.mu.Lock()
	defer tp.mu.Unlock()
	return tp.actual[net.JoinHostPort("192.0.2.1", strconv.Itoa(int(port)))]
}

// named is a backend that answers with its name and the request it saw.
func named(t *testing.T, name string) (*httptest.Server, Endpoint) {
	s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("X-Backend", name)
		n, _ := io.Copy(io.Discard, r.Body)
		fmt.Fprintf(w, "%s path=%s host=%s xff=%s proto=%s xhost=%s realip=%s added=%s removed=%s body=%d",
			name, r.URL.Path, r.Host, r.Header.Get("X-Forwarded-For"), r.Header.Get("X-Forwarded-Proto"),
			r.Header.Get("X-Forwarded-Host"), r.Header.Get("X-Real-IP"), strings.Join(r.Header.Values("X-Added"), ","),
			r.Header.Get("X-Remove-Me"), n)
	}))
	t.Cleanup(s.Close)
	_, port, _ := net.SplitHostPort(s.Listener.Addr().String())
	p, _ := strconv.Atoi(port)
	return s, Endpoint{IP: net.ParseIP("127.0.0.1").To4(), Port: int32(p)}
}

func backend(name string, ep Endpoint) []ServeBackend {
	return []ServeBackend{{Key: name, Weight: 1, Endpoints: []Endpoint{ep}}}
}

func prefix(v string) HTTPRouteMatch {
	return HTTPRouteMatch{Path: &HTTPPathMatch{Type: ptr(pathPrefix), Value: ptr(v)}}
}

func exact(v string) HTTPRouteMatch {
	return HTTPRouteMatch{Path: &HTTPPathMatch{Type: ptr(pathExact), Value: ptr(v)}}
}

func get(t *testing.T, url, host string, hdr map[string]string) (int, string, http.Header) {
	t.Helper()
	req, _ := http.NewRequest("GET", url, nil)
	if host != "" {
		req.Host = host
	}
	for k, v := range hdr {
		req.Header.Set(k, v)
	}
	c := &http.Client{CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}
	resp, err := c.Do(req)
	if err != nil {
		t.Fatalf("GET %s: %v", url, err)
	}
	defer resp.Body.Close()
	b, _ := io.ReadAll(resp.Body)
	return resp.StatusCode, string(b), resp.Header
}

func TestProxyRoutesByPrecedence(t *testing.T) {
	tp := newTestProxy(t)
	_, a := named(t, "A")
	_, b := named(t, "B")
	_, c := named(t, "C")
	_, d := named(t, "D")
	_, e := named(t, "E")
	t0 := time.Date(2026, 10, 14, 1, 0, 0, 0, time.UTC)
	routes := []ServeRoute{
		{Key: "ns/first", Created: t0, Rules: []ServeRule{
			{Matches: []HTTPRouteMatch{exact("/exact")}, Backends: backend("A", a)},
			{Matches: []HTTPRouteMatch{prefix("/api")}, Backends: backend("B", b)},
			{Matches: []HTTPRouteMatch{prefix("/")}, Backends: backend("C", c)},
		}},
		{Key: "ns/second", Created: t0.Add(time.Minute), Rules: []ServeRule{
			{Matches: []HTTPRouteMatch{prefix("/api/v2")}, Backends: backend("D", d)},
			{Matches: []HTTPRouteMatch{{Path: &HTTPPathMatch{Type: ptr(pathPrefix), Value: ptr("/api")},
				Headers: []HTTPHeaderMatch{{Name: "X-Ver", Value: "2"}}}}, Backends: backend("E", e)},
		}},
	}
	if f := tp.Apply([]ServeListener{{Gateway: "ns/gw", Name: "http", Address: net.ParseIP("192.0.2.1"), Port: 8080, Routes: routes}}); len(f) != 0 {
		t.Fatal(f)
	}
	base := "http://" + tp.at(8080)
	for path, want := range map[string]string{
		"/exact": "A", "/exactly": "C", "/api/x": "B", "/api": "B", "/api/v2/y": "D", "/apix": "C", "/": "C",
	} {
		if _, body, _ := get(t, base+path, "example.com", nil); !strings.HasPrefix(body, want+" ") {
			t.Errorf("%s went to %q, want %s", path, body, want)
		}
	}
	// same prefix length: the match with more headers takes it, though its route is younger
	if _, body, _ := get(t, base+"/api/x", "example.com", map[string]string{"X-Ver": "2"}); !strings.HasPrefix(body, "E ") {
		t.Errorf("header match lost: %q", body)
	}
}

func TestProxyHostPrecedence(t *testing.T) {
	tp := newTestProxy(t)
	_, w := named(t, "W")
	_, x := named(t, "X")
	routes := []ServeRoute{
		{Key: "ns/wild", Hostnames: []string{"*.example.com"}, Rules: []ServeRule{{Matches: []HTTPRouteMatch{prefix("/")}, Backends: backend("W", w)}}},
		{Key: "ns/exact", Hostnames: []string{"a.example.com"}, Rules: []ServeRule{{Matches: []HTTPRouteMatch{prefix("/")}, Backends: backend("X", x)}}},
	}
	tp.Apply([]ServeListener{{Gateway: "ns/gw", Name: "http", Address: net.ParseIP("192.0.2.1"), Port: 8081, Routes: routes}})
	base := "http://" + tp.at(8081)
	for host, want := range map[string]string{"a.example.com": "X", "A.Example.com:8081": "X", "b.example.com": "W"} {
		if _, body, _ := get(t, base+"/", host, nil); !strings.HasPrefix(body, want+" ") {
			t.Errorf("host %s went to %q, want %s", host, body, want)
		}
	}
	if code, _, _ := get(t, base+"/", "example.org", nil); code != http.StatusNotFound {
		t.Errorf("a host no route names: %d, want 404", code)
	}
}

func certSecret(t *testing.T, name, rv string) *coreapi.Secret {
	s := tlsSecret(t, "gateway-system", strings.ReplaceAll(name, ".", "-"), name)
	s.UID = types.UID("uid-" + name)
	s.ResourceVersion = rv
	return s
}

func dialTLS(t *testing.T, addr, sni string) (*x509.Certificate, error) {
	conn, err := tls.Dial("tcp", addr, &tls.Config{ServerName: sni, InsecureSkipVerify: true})
	if err != nil {
		return nil, err
	}
	defer conn.Close()
	return conn.ConnectionState().PeerCertificates[0], nil
}

func TestProxyTLSBySNIAndSwapsCertificates(t *testing.T) {
	tp := newTestProxy(t)
	_, a := named(t, "A")
	_, b := named(t, "B")
	listeners := func(certA *coreapi.Secret) []ServeListener {
		return []ServeListener{
			{Gateway: "ns/gw", Name: "a", Address: net.ParseIP("192.0.2.1"), Port: 8443, HTTPS: true, Hostname: "a.test", Cert: certA,
				Routes: []ServeRoute{{Key: "ns/a", Rules: []ServeRule{{Matches: []HTTPRouteMatch{prefix("/")}, Backends: backend("A", a)}}}}},
			{Gateway: "ns/gw", Name: "b", Address: net.ParseIP("192.0.2.1"), Port: 8443, HTTPS: true, Hostname: "b.test", Cert: certSecret(t, "b.test", "1"),
				Routes: []ServeRoute{{Key: "ns/b", Rules: []ServeRule{{Matches: []HTTPRouteMatch{prefix("/")}, Backends: backend("B", b)}}}}},
		}
	}
	first := certSecret(t, "a.test", "1")
	if f := tp.Apply(listeners(first)); len(f) != 0 {
		t.Fatal(f)
	}
	addr := tp.at(8443)
	ca, err := dialTLS(t, addr, "a.test")
	if err != nil || ca.Subject.CommonName != "a.test" {
		t.Fatalf("SNI a.test: %v %v", ca, err)
	}
	cb, err := dialTLS(t, addr, "b.test")
	if err != nil || cb.Subject.CommonName != "b.test" {
		t.Fatalf("SNI b.test: %v %v", cb, err)
	}
	if _, err := dialTLS(t, addr, "c.test"); err == nil {
		t.Fatal("a name no listener has must not get a certificate")
	}
	// requests go by listener, over HTTP/2 when the client offers it
	client := &http.Client{Transport: &http.Transport{TLSClientConfig: &tls.Config{InsecureSkipVerify: true, ServerName: "b.test"}, ForceAttemptHTTP2: true}}
	req, _ := http.NewRequest("GET", "https://"+addr+"/", nil)
	req.Host = "b.test"
	resp, err := client.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	body, _ := io.ReadAll(resp.Body)
	resp.Body.Close()
	if !strings.HasPrefix(string(body), "B ") || !strings.Contains(string(body), "proto=https") || resp.ProtoMajor != 2 {
		t.Fatalf("over HTTPS to b.test: %s %q", resp.Proto, body)
	}
	// a renewed Secret: new handshakes get the new certificate on the same, still-bound port
	renewed := certSecret(t, "a.test", "2")
	tp.Apply(listeners(renewed))
	if tp.at(8443) != addr {
		t.Fatal("the port was rebound for a certificate change")
	}
	cn, err := dialTLS(t, addr, "a.test")
	if err != nil || cn.SerialNumber.Cmp(ca.SerialNumber) == 0 {
		t.Fatalf("still the old certificate after the swap: %v", err)
	}
}

func TestProxyRedirectsAndForwards(t *testing.T) {
	tp := newTestProxy(t)
	_, a := named(t, "A")
	routes := []ServeRoute{{Key: "ns/r", Rules: []ServeRule{
		{Matches: []HTTPRouteMatch{prefix("/old")}, Filters: []HTTPRouteFilter{{Type: filterRequestRedirect,
			RequestRedirect: &HTTPRequestRedirectFilter{Scheme: ptr("https"), StatusCode: ptr(301)}}}},
		{Matches: []HTTPRouteMatch{prefix("/moved")}, Filters: []HTTPRouteFilter{{Type: filterRequestRedirect,
			RequestRedirect: &HTTPRequestRedirectFilter{Path: &HTTPPathModifier{Type: "ReplacePrefixMatch", ReplacePrefixMatch: ptr("/new")}}}}},
		{Matches: []HTTPRouteMatch{prefix("/v1")}, Backends: backend("A", a), Filters: []HTTPRouteFilter{
			{Type: filterRequestHeaderModifier, RequestHeaderModifier: &HTTPHeaderFilter{
				Set: []HTTPHeader{{Name: "X-Set", Value: "s"}}, Add: []HTTPHeader{{Name: "X-Added", Value: "two"}}, Remove: []string{"X-Remove-Me"}}},
			{Type: filterResponseHeaderModifier, ResponseHeaderModifier: &HTTPHeaderFilter{Set: []HTTPHeader{{Name: "X-Resp", Value: "r"}}}},
			{Type: filterURLRewrite, URLRewrite: &HTTPURLRewriteFilter{Path: &HTTPPathModifier{Type: "ReplacePrefixMatch", ReplacePrefixMatch: ptr("/v2")}}},
		}},
		{Matches: []HTTPRouteMatch{prefix("/")}, Backends: backend("A", a)},
	}}}
	tp.Apply([]ServeListener{{Gateway: "ns/gw", Name: "http", Address: net.ParseIP("192.0.2.1"), Port: 80, Routes: routes}})
	base := "http://" + tp.at(80)
	code, _, h := get(t, base+"/old/page?q=1", "git.levk.dev", nil)
	if code != http.StatusMovedPermanently || h.Get("Location") != "https://git.levk.dev/old/page?q=1" {
		t.Fatalf("scheme redirect: %d %q", code, h.Get("Location"))
	}
	code, _, h = get(t, base+"/moved/x", "git.levk.dev", nil)
	if code != http.StatusFound || h.Get("Location") != "http://git.levk.dev/new/x" {
		t.Fatalf("prefix redirect: %d %q", code, h.Get("Location"))
	}
	// the client's own claims about who it is are not passed on; the backend sees this hop's view
	_, body, _ := get(t, base+"/", "git.levk.dev", map[string]string{"X-Forwarded-For": "6.6.6.6", "X-Real-IP": "6.6.6.6"})
	if !strings.Contains(body, "host=git.levk.dev") || !strings.Contains(body, "xff=127.0.0.1 ") ||
		!strings.Contains(body, "proto=http ") || !strings.Contains(body, "xhost=git.levk.dev") || !strings.Contains(body, "realip=127.0.0.1") {
		t.Fatalf("forwarded headers: %q", body)
	}
	_, body, h = get(t, base+"/v1/thing", "git.levk.dev", map[string]string{"X-Added": "one", "X-Remove-Me": "x"})
	if !strings.Contains(body, "path=/v2/thing") || !strings.Contains(body, "added=one,two") || !strings.Contains(body, "removed= ") || h.Get("X-Resp") != "r" {
		t.Fatalf("filters: %q %v", body, h)
	}
}

func TestProxyAnswersWhatItCannotRoute(t *testing.T) {
	tp := newTestProxy(t)
	routes := []ServeRoute{{Key: "ns/r", Rules: []ServeRule{
		{Matches: []HTTPRouteMatch{prefix("/broken")}, Broken: true},
		{Matches: []HTTPRouteMatch{prefix("/empty")}, Backends: []ServeBackend{{Key: "ns/svc:80", Weight: 1}}},
	}}}
	tp.Apply([]ServeListener{{Gateway: "ns/gw", Name: "http", Address: net.ParseIP("192.0.2.1"), Port: 8082, Routes: routes}})
	base := "http://" + tp.at(8082)
	for path, want := range map[string]int{"/broken": 500, "/empty": 503, "/nothing": 404} {
		if code, _, _ := get(t, base+path, "x", nil); code != want {
			t.Errorf("%s: %d, want %d", path, code, want)
		}
	}
}

func TestProxyStreamsLargeBodies(t *testing.T) {
	tp := newTestProxy(t)
	_, a := named(t, "A")
	tp.Apply([]ServeListener{{Gateway: "ns/gw", Name: "http", Address: net.ParseIP("192.0.2.1"), Port: 8083,
		Routes: []ServeRoute{{Key: "ns/r", Rules: []ServeRule{{Matches: []HTTPRouteMatch{prefix("/")}, Backends: backend("A", a)}}}}}})
	const size = 32 << 20
	resp, err := http.Post("http://"+tp.at(8083)+"/upload", "application/octet-stream", bytes.NewReader(make([]byte, size)))
	if err != nil {
		t.Fatal(err)
	}
	body, _ := io.ReadAll(resp.Body)
	resp.Body.Close()
	if !strings.Contains(string(body), fmt.Sprintf("body=%d", size)) {
		t.Fatalf("upload: %q", body)
	}
}

func TestProxyBindFailureAndClose(t *testing.T) {
	tp := newTestProxy(t)
	tp.listen = func(addr string) (net.Listener, error) { return nil, fmt.Errorf("address in use") }
	f := tp.Apply([]ServeListener{{Gateway: "ns/gw", Name: "http", Address: net.ParseIP("192.0.2.1"), Port: 8084}})
	if f["ns/gw/http"] == nil {
		t.Fatalf("bind failure not reported: %v", f)
	}
	tp2 := newTestProxy(t)
	tp2.Apply([]ServeListener{{Gateway: "ns/gw", Name: "http", Address: net.ParseIP("192.0.2.1"), Port: 8085}})
	addr := tp2.at(8085)
	tp2.Apply(nil)
	deadline := time.Now().Add(5 * time.Second)
	for time.Now().Before(deadline) {
		c, err := net.DialTimeout("tcp", addr, 200*time.Millisecond)
		if err != nil {
			return
		}
		c.Close()
		time.Sleep(50 * time.Millisecond)
	}
	t.Fatal("a port no longer served still accepts connections")
}