package host

import (
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"
	"testing/fstest"
)

// upstream records what reached the homeserver.
type upstream struct {
	srv  *httptest.Server
	hits []*http.Request
}

func newUpstream(t *testing.T) *upstream {
	t.Helper()
	u := &upstream{}
	u.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		u.hits = append(u.hits, r.Clone(r.Context()))
		w.Header().Set("Content-Type", "application/json")
		w.Header().Set("Set-Cookie", "tracker=1")
		_, _ = w.Write([]byte(`{"path":"` + r.URL.EscapedPath() + `"}`))
	}))
	t.Cleanup(u.srv.Close)
	return u
}

var panelFS = fstest.MapFS{
	"index.html":         {Data: []byte("<!doctype html><title>panel</title>")},
	"favicon.svg":        {Data: []byte("<svg/>")},
	"assets/index-a1.js": {Data: []byte("console.log(1)")},
}

func newHandler(t *testing.T, up *upstream, panel bool) http.Handler {
	t.Helper()
	target, _ := url.Parse(up.srv.URL)
	opts := Options{Homeserver: target}
	if panel {
		opts.Panel = panelFS
	}
	h, err := New(opts)
	if err != nil {
		t.Fatal(err)
	}
	return h
}

func do(h http.Handler, method, target string, header http.Header) *httptest.ResponseRecorder {
	req := httptest.NewRequest(method, target, nil)
	for k, v := range header {
		req.Header[k] = v
	}
	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, req)
	return rec
}

func TestProxiesAdminAndSessionPaths(t *testing.T) {
	up := newUpstream(t)
	h := newHandler(t, up, true)

	for _, p := range []string{
		"/_matrix/client/versions",
		"/_matrix/client/v3/login",
		"/_matrix/client/v3/logout",
		"/_matrix/client/v3/account/whoami",
		"/_warp/admin/v1/stats",
		"/_warp/admin/v1/users/%40alice%3Awarp.example/devices",
		"/_warp/admin/v1/audit_log?limit=5&action=user.create",
	} {
		rec := do(h, http.MethodGet, p, nil)
		if rec.Code != http.StatusOK {
			t.Errorf("%s: status %d, want 200", p, rec.Code)
		}
	}
	if len(up.hits) != 7 {
		t.Fatalf("upstream saw %d requests, want 7", len(up.hits))
	}
	// Escaped IDs and the query reach the homeserver unchanged.
	if got := up.hits[5].URL.EscapedPath(); got != "/_warp/admin/v1/users/%40alice%3Awarp.example/devices" {
		t.Errorf("escaped path = %q", got)
	}
	if got := up.hits[6].URL.RawQuery; got != "limit=5&action=user.create" {
		t.Errorf("query = %q", got)
	}
}

func TestRefusesOtherHomeserverPaths(t *testing.T) {
	up := newUpstream(t)
	h := newHandler(t, up, true)

	for _, p := range []string{
		"/_matrix/client/v3/sync",
		"/_matrix/client/v3/rooms/!a:b/send/m.room.message/1",
		"/_matrix/federation/v1/version",
		"/_matrix/media/v3/upload",
		"/_warp/admin/v2/stats",
		"/_warp/admin/v1/../../_matrix/client/v3/sync",
		"/_warp/admin/v1/%2e%2e/%2e%2e/%2e%2e/_matrix/client/v3/sync",
		"/_warp/admin/v1//stats",
	} {
		rec := do(h, http.MethodGet, p, nil)
		if rec.Code != http.StatusNotFound {
			t.Errorf("%s: status %d, want 404", p, rec.Code)
		}
		var body map[string]string
		_ = json.Unmarshal(rec.Body.Bytes(), &body)
		if body["errcode"] != "M_UNRECOGNIZED" {
			t.Errorf("%s: errcode %q", p, body["errcode"])
		}
	}
	if len(up.hits) != 0 {
		t.Fatalf("upstream saw %d requests, want none", len(up.hits))
	}
}

func TestProxyStripsCookiesAndForwardsToken(t *testing.T) {
	up := newUpstream(t)
	h := newHandler(t, up, true)

	rec := do(h, http.MethodGet, "/_warp/admin/v1/stats", http.Header{
		"Authorization": {"Bearer tok"},
		"Cookie":        {"session=secret"},
	})
	if rec.Header().Get("Set-Cookie") != "" {
		t.Error("Set-Cookie from the homeserver was passed to the browser")
	}
	if rec.Header().Get("Cache-Control") != "no-store" {
		t.Errorf("Cache-Control = %q, want no-store", rec.Header().Get("Cache-Control"))
	}
	got := up.hits[0]
	if got.Header.Get("Cookie") != "" {
		t.Error("Cookie was passed to the homeserver")
	}
	if got.Header.Get("Authorization") != "Bearer tok" {
		t.Errorf("Authorization = %q", got.Header.Get("Authorization"))
	}
	if got.Header.Get("X-Forwarded-For") == "" {
		t.Error("X-Forwarded-For not set")
	}
}

func TestUnreachableHomeserverIs502(t *testing.T) {
	target, _ := url.Parse("http://127.0.0.1:1")
	h, err := New(Options{Homeserver: target, Panel: panelFS})
	if err != nil {
		t.Fatal(err)
	}
	rec := do(h, http.MethodGet, "/_warp/admin/v1/stats", nil)
	if rec.Code != http.StatusBadGateway {
		t.Fatalf("status %d, want 502", rec.Code)
	}
	if !strings.Contains(rec.Body.String(), "M_UNKNOWN") {
		t.Errorf("body %q", rec.Body.String())
	}
}

func TestServesPanel(t *testing.T) {
	up := newUpstream(t)
	h := newHandler(t, up, true)

	cases := []struct {
		path, wantBody, wantCache string
		wantCode                  int
	}{
		{"/", "panel", "no-cache", 200},
		{"/users/%40alice%3Awarp.example", "panel", "no-cache", 200},
		{"/rooms/!abc:warp.example", "panel", "no-cache", 200},
		{"/index.html", "panel", "no-cache", 200},
		{"/assets/index-a1.js", "console.log", "public, max-age=31536000, immutable", 200},
		{"/favicon.svg", "<svg", "public, max-age=3600", 200},
		{"/assets/gone-b2.js", "", "", 404},
		{"/robots.txt", "", "", 404},
	}
	for _, c := range cases {
		rec := do(h, http.MethodGet, c.path, nil)
		if rec.Code != c.wantCode {
			t.Errorf("%s: status %d, want %d", c.path, rec.Code, c.wantCode)
			continue
		}
		if c.wantCode != 200 {
			continue
		}
		body, _ := io.ReadAll(rec.Body)
		if !strings.Contains(string(body), c.wantBody) {
			t.Errorf("%s: body %q", c.path, body)
		}
		if got := rec.Header().Get("Cache-Control"); got != c.wantCache {
			t.Errorf("%s: Cache-Control %q, want %q", c.path, got, c.wantCache)
		}
		if rec.Header().Get("Content-Security-Policy") == "" {
			t.Errorf("%s: no CSP", c.path)
		}
		if rec.Header().Get("X-Frame-Options") != "DENY" {
			t.Errorf("%s: X-Frame-Options %q", c.path, rec.Header().Get("X-Frame-Options"))
		}
	}

	if rec := do(h, http.MethodPost, "/", nil); rec.Code != http.StatusMethodNotAllowed {
		t.Errorf("POST /: status %d, want 405", rec.Code)
	}
}

func TestNotBuiltPanel(t *testing.T) {
	up := newUpstream(t)
	h := newHandler(t, up, false)

	rec := do(h, http.MethodGet, "/", nil)
	if rec.Code != http.StatusServiceUnavailable || !strings.Contains(rec.Body.String(), "isn’t built") {
		t.Fatalf("status %d body %q", rec.Code, rec.Body.String())
	}
	// The proxy still works without the panel.
	if rec := do(h, http.MethodGet, "/_warp/admin/v1/stats", nil); rec.Code != http.StatusOK {
		t.Errorf("proxy status %d", rec.Code)
	}
}

func TestHealthz(t *testing.T) {
	up := newUpstream(t)
	h := newHandler(t, up, true)
	if rec := do(h, http.MethodGet, "/healthz", nil); rec.Code != 200 || rec.Body.String() != "ok\n" {
		t.Fatalf("status %d body %q", rec.Code, rec.Body.String())
	}
}

func TestNewRejectsRelativeHomeserver(t *testing.T) {
	u, _ := url.Parse("warp:8008")
	if _, err := New(Options{Homeserver: u}); err == nil {
		t.Fatal("want an error for a URL without a scheme and host")
	}
}