package tunnel

import (
	"context"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/coder/websocket"
)

// The whole loop, no router required: a local "instance" (plain HTTP +
// a WebSocket echo), a gateway, and a client tunnelling between them. Verifies
// ordinary requests, header passthrough, and — the part that matters for the
// chat — a live WebSocket upgraded end-to-end through the tunnel.
func TestTunnelEndToEnd(t *testing.T) {
	backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/api/version":
			w.Header().Set("Content-Type", "application/json")
			io.WriteString(w, `{"build":"test","version":"test-1"}`)
		case "/ws":
			c, err := websocket.Accept(w, r, nil)
			if err != nil {
				return
			}
			defer c.CloseNow()
			ctx := r.Context()
			for {
				typ, data, err := c.Read(ctx)
				if err != nil {
					return
				}
				if err := c.Write(ctx, typ, append([]byte("echo:"), data...)); err != nil {
					return
				}
			}
		default:
			w.Header().Set("X-Seen-Host", r.Host)
			io.WriteString(w, "hello through the tunnel")
		}
	}))
	defer backend.Close()

	gw := &Gateway{Domain: "test.local", DialTimeout: 5 * time.Second}
	registered := make(chan string, 1)
	gw.OnRegister = func(host string) error { registered <- host; return nil }
	closed := make(chan string, 1)
	gw.OnClose = func(host string) { closed <- host }

	gwSrv := httptest.NewServer(gw)
	defer gwSrv.Close()
	u, _ := url.Parse(gwSrv.URL)
	gw.ControlHost = u.Hostname() // ServeHTTP compares hosts port-stripped

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	cl := &Client{Gateway: gwSrv.URL, LocalAddr: strings.TrimPrefix(backend.URL, "http://")}
	ready := make(chan string, 1)
	clientDone := make(chan error, 1)
	go func() { clientDone <- cl.Run(ctx, func(h string) { ready <- h }) }()

	var host string
	select {
	case host = <-ready:
	case err := <-clientDone:
		t.Fatalf("client exited early: %v", err)
	case <-time.After(5 * time.Second):
		t.Fatal("no ready")
	}
	if !strings.HasPrefix(host, "try-") || !strings.HasSuffix(host, ".test.local") {
		t.Fatalf("unexpected host %q", host)
	}
	if got := <-registered; got != host {
		t.Fatalf("OnRegister got %q, ready got %q", got, host)
	}

	// A plain request, addressed to the slug host through the gateway.
	hostRewrite := func(rt http.RoundTripper) http.RoundTripper {
		return roundTripperFunc(func(r *http.Request) (*http.Response, error) {
			r.Host = host
			return rt.RoundTrip(r)
		})
	}
	httpc := &http.Client{Transport: hostRewrite(http.DefaultTransport.(*http.Transport).Clone())}
	resp, err := httpc.Get(gwSrv.URL + "/anything")
	if err != nil {
		t.Fatalf("GET: %v", err)
	}
	body, _ := io.ReadAll(resp.Body)
	resp.Body.Close()
	if string(body) != "hello through the tunnel" {
		t.Fatalf("body = %q", body)
	}
	if got := resp.Header.Get("X-Seen-Host"); got != host {
		t.Fatalf("backend saw Host %q, want %q (public hostname must reach the instance)", got, host)
	}

	// A WebSocket, upgraded end-to-end through the tunnel.
	wsCtx, wsCancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer wsCancel()
	c, _, err := websocket.Dial(wsCtx, "ws://"+u.Host+"/ws", &websocket.DialOptions{
		HTTPClient: &http.Client{Transport: hostRewrite(http.DefaultTransport.(*http.Transport).Clone())},
	})
	if err != nil {
		t.Fatalf("ws dial: %v", err)
	}
	defer c.CloseNow()
	if err := c.Write(wsCtx, websocket.MessageText, []byte("ping")); err != nil {
		t.Fatalf("ws write: %v", err)
	}
	_, data, err := c.Read(wsCtx)
	if err != nil {
		t.Fatalf("ws read: %v", err)
	}
	if string(data) != "echo:ping" {
		t.Fatalf("ws echo = %q", data)
	}

	// Tearing the client down retires the route.
	cancel()
	select {
	case got := <-closed:
		if got != host {
			t.Fatalf("OnClose got %q", got)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("route never closed")
	}
}

// The per-IP limiter refuses the sixth tunnel in an hour from one address.
func TestTunnelRateLimit(t *testing.T) {
	gw := &Gateway{Domain: "test.local"}
	for range 5 {
		if err := gw.admit("203.0.113.9"); err != nil {
			t.Fatalf("admit: %v", err)
		}
		gw.byIP["203.0.113.9"] = append(gw.byIP["203.0.113.9"], time.Now())
	}
	if err := gw.admit("203.0.113.9"); err == nil {
		t.Fatal("sixth tunnel admitted")
	}
	if err := gw.admit("203.0.113.10"); err != nil {
		t.Fatalf("other address blocked: %v", err)
	}
}

// A backend that is not an Eleven instance is refused before its hostname is
// ever published.
func TestTunnelRejectsNonEleven(t *testing.T) {
	backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		io.WriteString(w, "just a web server")
	}))
	defer backend.Close()

	gw := &Gateway{Domain: "test.local", DialTimeout: 5 * time.Second}
	registered := false
	gw.OnRegister = func(string) error { registered = true; return nil }
	gwSrv := httptest.NewServer(gw)
	defer gwSrv.Close()
	u, _ := url.Parse(gwSrv.URL)
	gw.ControlHost = u.Hostname()

	cl := &Client{Gateway: gwSrv.URL, LocalAddr: strings.TrimPrefix(backend.URL, "http://")}
	err := cl.Run(context.Background(), func(string) {})
	if err == nil || !strings.Contains(err.Error(), "only Eleven instances") {
		t.Fatalf("expected probe refusal, got %v", err)
	}
	if registered {
		t.Fatal("route was published for a non-Eleven backend")
	}
}

// The per-tunnel connection cap refuses dial-backs once live + in-flight
// connections hit the limit.
func TestTunnelConnCap(t *testing.T) {
	g := &Gateway{MaxConns: 2}
	tun := &liveTunnel{closed: make(chan struct{})}
	tun.conns.Store(2)
	if _, err := g.dialBack(context.Background(), tun); err == nil ||
		!strings.Contains(err.Error(), "connection cap") {
		t.Fatalf("expected cap refusal, got %v", err)
	}
	if got := tun.conns.Load(); got != 2 {
		t.Fatalf("refused dial leaked a slot: %d", got)
	}
}

// startTunnelPair boots a backend serving handler, the given gateway, and a
// client tunnelling between them, returning the gateway server and the public
// slug host once it's live. Everything tears down with the test.
func startTunnelPair(t *testing.T, gw *Gateway, handler http.HandlerFunc) (gwSrv *httptest.Server, host string) {
	t.Helper()
	backend := httptest.NewServer(handler)
	t.Cleanup(backend.Close)
	gwSrv = httptest.NewServer(gw)
	t.Cleanup(gwSrv.Close)
	u, _ := url.Parse(gwSrv.URL)
	gw.ControlHost = u.Hostname()

	ctx, cancel := context.WithCancel(context.Background())
	t.Cleanup(cancel)
	cl := &Client{Gateway: gwSrv.URL, LocalAddr: strings.TrimPrefix(backend.URL, "http://")}
	ready := make(chan string, 1)
	clientDone := make(chan error, 1)
	go func() { clientDone <- cl.Run(ctx, func(h string) { ready <- h }) }()
	select {
	case host = <-ready:
	case err := <-clientDone:
		t.Fatalf("client exited early: %v", err)
	case <-time.After(5 * time.Second):
		t.Fatal("no ready")
	}
	return gwSrv, host
}

// hostRewriteTransport addresses requests to the slug host through the
// gateway's test listener.
func hostRewriteTransport(host string) http.RoundTripper {
	rt := http.DefaultTransport.(*http.Transport).Clone()
	return roundTripperFunc(func(r *http.Request) (*http.Response, error) {
		r.Host = host
		return rt.RoundTrip(r)
	})
}

// versionThen serves the probe's /api/version and hands everything else to
// next — the minimum for a backend to pass the Eleven door check.
func versionThen(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/api/version" {
			w.Header().Set("Content-Type", "application/json")
			io.WriteString(w, `{"build":"test","version":"test-1"}`)
			return
		}
		next(w, r)
	}
}

// A first-load burst wider than the per-tunnel connection cap queues for a
// free connection instead of failing. The no-build shell fetches 40+ modules
// in one burst over HTTP/2, and a single 502 permanently fails the page's ES
// module graph — this is what made `lchat try` unusable in a browser.
func TestTunnelBurstQueuesBeyondConnCap(t *testing.T) {
	gw := &Gateway{Domain: "test.local", DialTimeout: 5 * time.Second, MaxConns: 4}
	gwSrv, host := startTunnelPair(t, gw, versionThen(func(w http.ResponseWriter, r *http.Request) {
		time.Sleep(20 * time.Millisecond) // hold connections open so the burst overlaps
		io.WriteString(w, "module")
	}))

	httpc := &http.Client{Transport: hostRewriteTransport(host)}
	var wg sync.WaitGroup
	codes := make([]int, 24)
	for i := range codes {
		wg.Add(1)
		go func() {
			defer wg.Done()
			resp, err := httpc.Get(fmt.Sprintf("%s/static/m%d.js", gwSrv.URL, i))
			if err != nil {
				codes[i] = -1
				return
			}
			io.Copy(io.Discard, resp.Body)
			resp.Body.Close()
			codes[i] = resp.StatusCode
		}()
	}
	wg.Wait()
	for i, c := range codes {
		if c != http.StatusOK {
			t.Fatalf("request %d: status %d — a burst above the cap must queue, not fail", i, c)
		}
	}
}

// A visitor abandoning a request mid-dial (reload, closed tab) fails that one
// dial only — it must never drop the tunnel for everyone else.
func TestTunnelSurvivesAbandonedRequest(t *testing.T) {
	closed := make(chan string, 1)
	gw := &Gateway{Domain: "test.local", DialTimeout: 5 * time.Second}
	gw.OnClose = func(host string) { closed <- host }
	gwSrv, host := startTunnelPair(t, gw, versionThen(func(w http.ResponseWriter, r *http.Request) {
		io.WriteString(w, "ok")
	}))

	gw.mu.Lock()
	tun := gw.tunnels[host]
	gw.mu.Unlock()
	if tun == nil {
		t.Fatal("no live tunnel")
	}

	abandoned, cancel := context.WithCancel(context.Background())
	cancel()
	if conn, err := gw.dialBack(abandoned, tun); err == nil {
		conn.Close() // the dial-back can win the race; the tunnel must survive either way
	}

	select {
	case h := <-closed:
		t.Fatalf("abandoned request dropped the whole tunnel (%s)", h)
	default:
	}
	httpc := &http.Client{Transport: hostRewriteTransport(host)}
	resp, err := httpc.Get(gwSrv.URL + "/after")
	if err != nil {
		t.Fatalf("tunnel unusable after an abandoned request: %v", err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("tunnel unusable after an abandoned request: HTTP %d", resp.StatusCode)
	}
}

type roundTripperFunc func(*http.Request) (*http.Response, error)

func (f roundTripperFunc) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) }