package main

import (
	"context"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/coder/websocket"
)

func TestRouterProxiesByHost(t *testing.T) {
	// A backend on a unix socket that echoes the Host header it receives.
	sock := tmpSock(t)
	ln, err := net.Listen("unix", sock)
	if err != nil {
		t.Fatalf("listen unix: %v", err)
	}
	backend := &http.Server{Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		io.WriteString(w, "backend:"+r.Host)
	})}
	go backend.Serve(ln)
	t.Cleanup(func() { backend.Close() })

	reg, err := openRegistry(filepath.Join(t.TempDir(), "registry.db"))
	if err != nil {
		t.Fatalf("open registry: %v", err)
	}
	if err := reg.upsert(Route{Host: "chat.example.com", Kind: "instance", Socket: sock}); err != nil {
		t.Fatalf("upsert: %v", err)
	}
	rt, err := newRouter(reg, nil)
	if err != nil {
		t.Fatalf("new router: %v", err)
	}
	front := httptest.NewServer(rt)
	t.Cleanup(front.Close)

	// Known host is proxied, and the backend sees the real Host header.
	if body, code := proxyGet(t, front.URL, "chat.example.com"); code != 200 || body != "backend:chat.example.com" {
		t.Fatalf("known host: got %d %q, want 200 backend:chat.example.com", code, body)
	}
	// Unknown host → 404 (and, as HostPolicy, no cert would be issued for it).
	if _, code := proxyGet(t, front.URL, "nope.example.com"); code != http.StatusNotFound {
		t.Fatalf("unknown host: got %d, want 404", code)
	}
	if rt.allowed("nope.example.com") {
		t.Fatal("allowed() should be false for an unconfigured host")
	}

	// A new route isn't served until reload; after reload it is (atomic swap).
	if err := reg.upsert(Route{Host: "two.example.com", Kind: "instance", Socket: sock}); err != nil {
		t.Fatalf("upsert two: %v", err)
	}
	if rt.allowed("two.example.com") {
		t.Fatal("route should not be live before reload")
	}
	if err := rt.reload(); err != nil {
		t.Fatalf("reload: %v", err)
	}
	if body, code := proxyGet(t, front.URL, "two.example.com"); code != 200 || body != "backend:two.example.com" {
		t.Fatalf("after reload: got %d %q, want 200 backend:two.example.com", code, body)
	}
}

func TestRouterProxiesWebSocket(t *testing.T) {
	// A backend WebSocket echo server on a unix socket.
	sock := tmpSock(t)
	ln, err := net.Listen("unix", sock)
	if err != nil {
		t.Fatalf("listen unix: %v", err)
	}
	backend := &http.Server{Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		c, err := websocket.Accept(w, r, nil)
		if err != nil {
			return
		}
		defer c.Close(websocket.StatusNormalClosure, "")
		mt, data, err := c.Read(r.Context())
		if err != nil {
			return
		}
		_ = c.Write(r.Context(), mt, append([]byte("echo:"), data...))
	})}
	go backend.Serve(ln)
	t.Cleanup(func() { backend.Close() })

	reg, err := openRegistry(filepath.Join(t.TempDir(), "registry.db"))
	if err != nil {
		t.Fatalf("open registry: %v", err)
	}
	// Route the loopback host so the client reaches it without a Host override
	// (ServeHTTP strips the port, matching "127.0.0.1").
	if err := reg.upsert(Route{Host: "127.0.0.1", Kind: "instance", Socket: sock}); err != nil {
		t.Fatalf("upsert: %v", err)
	}
	rt, err := newRouter(reg, nil)
	if err != nil {
		t.Fatalf("new router: %v", err)
	}
	front := httptest.NewServer(rt)
	t.Cleanup(front.Close)

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	wsURL := "ws" + strings.TrimPrefix(front.URL, "http") + "/ws"
	c, _, err := websocket.Dial(ctx, wsURL, nil)
	if err != nil {
		t.Fatalf("ws dial through router: %v", err)
	}
	defer c.Close(websocket.StatusNormalClosure, "")
	if err := c.Write(ctx, websocket.MessageText, []byte("hi")); err != nil {
		t.Fatalf("ws write: %v", err)
	}
	_, data, err := c.Read(ctx)
	if err != nil {
		t.Fatalf("ws read: %v", err)
	}
	if string(data) != "echo:hi" {
		t.Fatalf("ws echo through router: got %q, want %q", data, "echo:hi")
	}
}

func TestIsRetryableDialErr(t *testing.T) {
	sock := tmpSock(t) // a path with no listener → ENOENT
	if _, err := net.Dial("unix", sock); err == nil {
		t.Fatal("expected a dial error for a missing socket")
	} else if !isRetryableDialErr(err) {
		t.Errorf("missing socket: isRetryableDialErr(%v) = false, want true", err)
	}
	if isRetryableDialErr(context.Canceled) {
		t.Error("context.Canceled should not be retryable")
	}
}

func TestRouterRetriesAcrossRestart(t *testing.T) {
	sock := tmpSock(t)
	reg, err := openRegistry(filepath.Join(t.TempDir(), "registry.db"))
	if err != nil {
		t.Fatal(err)
	}
	if err := reg.upsert(Route{Host: "chat.example.com", Kind: "instance", Socket: sock}); err != nil {
		t.Fatal(err)
	}
	rt, err := newRouter(reg, nil)
	if err != nil {
		t.Fatal(err)
	}
	front := httptest.NewServer(rt)
	t.Cleanup(front.Close)

	// Fire the request before any backend exists; the router must retry the dial.
	type res struct {
		code int
		body string
	}
	ch := make(chan res, 1)
	go func() {
		req, _ := http.NewRequest(http.MethodGet, front.URL+"/", nil)
		req.Host = "chat.example.com"
		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			ch <- res{0, "err: " + err.Error()}
			return
		}
		b, _ := io.ReadAll(resp.Body)
		resp.Body.Close()
		ch <- res{resp.StatusCode, string(b)}
	}()

	// Bind the backend after a gap, simulating the instance restart window.
	time.Sleep(300 * time.Millisecond)
	ln, err := net.Listen("unix", sock)
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	backend := &http.Server{Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		io.WriteString(w, "late-backend")
	})}
	go backend.Serve(ln)
	t.Cleanup(func() { backend.Close() })

	select {
	case r := <-ch:
		if r.code != 200 || r.body != "late-backend" {
			t.Fatalf("retry across gap: got %d %q, want 200 late-backend", r.code, r.body)
		}
	case <-time.After(10 * time.Second):
		t.Fatal("request never completed; router did not retry across the gap")
	}
}

func TestReloadOnlyServesReady(t *testing.T) {
	reg, err := openRegistry(filepath.Join(t.TempDir(), "registry.db"))
	if err != nil {
		t.Fatal(err)
	}
	for _, r := range []Route{
		{Host: "ready.example.com", Kind: "instance", Socket: "/run/lchat/a.sock", Status: "ready"},
		{Host: "failed.example.com", Kind: "instance", Socket: "/run/lchat/b.sock", Status: "failed"},
		{Host: "creating.example.com", Kind: "instance", Socket: "/run/lchat/c.sock", Status: "creating"},
	} {
		if err := reg.upsert(r); err != nil {
			t.Fatal(err)
		}
	}
	rt, err := newRouter(reg, nil)
	if err != nil {
		t.Fatal(err)
	}
	if !rt.allowed("ready.example.com") {
		t.Error("ready route should be served")
	}
	if rt.allowed("failed.example.com") {
		t.Error("failed route should not be served or cert-eligible")
	}
	if rt.allowed("creating.example.com") {
		t.Error("creating route should not be served")
	}
}

func TestWaitHealthy(t *testing.T) {
	serve := func(status int) string {
		sock := tmpSock(t)
		ln, err := net.Listen("unix", sock)
		if err != nil {
			t.Fatal(err)
		}
		srv := &http.Server{Handler: http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(status) })}
		go srv.Serve(ln)
		t.Cleanup(func() { srv.Close() })
		return sock
	}
	if err := waitHealthy(serve(200), 3*time.Second); err != nil {
		t.Errorf("200 backend: waitHealthy = %v, want nil", err)
	}
	if err := waitHealthy(serve(500), 1*time.Second); err == nil {
		t.Error("500 backend: waitHealthy = nil, want error (a broken binary must fail the deploy gate)")
	}
}

// proxyGet sends GET base+"/" with an overridden Host header and returns the
// response body + status code.
func proxyGet(t *testing.T, base, host string) (string, int) {
	t.Helper()
	req, err := http.NewRequest(http.MethodGet, base+"/", nil)
	if err != nil {
		t.Fatalf("new request: %v", err)
	}
	req.Host = host
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatalf("request: %v", err)
	}
	defer resp.Body.Close()
	b, _ := io.ReadAll(resp.Body)
	return string(b), resp.StatusCode
}

// tmpSock returns a short unix socket path (the sun_path limit is ~104 bytes, so
// t.TempDir() can be too long on macOS) and schedules its removal.
func tmpSock(t *testing.T) string {
	t.Helper()
	s := fmt.Sprintf("/tmp/lchat-rt-%d-%d.sock", os.Getpid(), time.Now().UnixNano()%1_000_000)
	os.Remove(s)
	t.Cleanup(func() { os.Remove(s) })
	return s
}

func TestRouterUnknownHostFallback(t *testing.T) {
	reg, err := openRegistry(filepath.Join(t.TempDir(), "registry.db"))
	if err != nil {
		t.Fatalf("open registry: %v", err)
	}
	rt, err := newRouter(reg, nil)
	if err != nil {
		t.Fatalf("new router: %v", err)
	}
	rt.fallbackURL = "https://elevenmessenger.com"
	front := httptest.NewServer(rt)
	t.Cleanup(front.Close)

	// Unknown host → 302 to the fallback, never a bare 404. 302 (not 301) so a
	// typo'd host isn't cached as permanently moved.
	req, _ := http.NewRequest(http.MethodGet, front.URL+"/some/path", nil)
	req.Host = "typo.elevenmessenger.com"
	client := &http.Client{CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}
	resp, err := client.Do(req)
	if err != nil {
		t.Fatalf("request: %v", err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusFound {
		t.Fatalf("status = %d, want 302", resp.StatusCode)
	}
	if loc := resp.Header.Get("Location"); loc != "https://elevenmessenger.com" {
		t.Fatalf("Location = %q, want the fallback", loc)
	}

	// Without a fallback the classic 404 stands (covered again for regressions).
	rt.fallbackURL = ""
	if _, code := proxyGet(t, front.URL, "typo.elevenmessenger.com"); code != http.StatusNotFound {
		t.Fatalf("no fallback: status = %d, want 404", code)
	}
}