package main

import (
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"testing/fstest"
)

func testFiles() fstest.MapFS {
	return fstest.MapFS{
		"get.sh":           {Data: []byte("echo base=@@BASE@@ and again @@BASE@@\n")},
		"setup.sh":         {Data: []byte("curl @@BASE@@/steps/ollama.sh\n")},
		"uninstall.sh":     {Data: []byte("echo bye\n")},
		"add-disk.sh":      {Data: []byte("echo disk @@BASE@@\n")},
		"install-agent.sh": {Data: []byte("echo agent\n")},
		"steps/ollama.sh":  {Data: []byte("echo ollama @@BASE@@\n")},
	}
}

func get(t *testing.T, mux *http.ServeMux, path string, mutate func(*http.Request)) *httptest.ResponseRecorder {
	req := httptest.NewRequest("GET", path, nil)
	if mutate != nil {
		mutate(req)
	}
	rec := httptest.NewRecorder()
	mux.ServeHTTP(rec, req)
	return rec
}

func TestScriptsGetTheCallersAddressNotAFixedDomain(t *testing.T) {
	mux := http.NewServeMux()
	registerInstallers(mux, testFiles(), "")
	rec := get(t, mux, "/get", func(r *http.Request) { r.Host = "control.example.org"; r.Header.Set("X-Forwarded-Proto", "https") })
	if rec.Code != 200 || !strings.Contains(rec.Body.String(), "base=https://control.example.org and again https://control.example.org") {
		t.Fatalf("%d %q", rec.Code, rec.Body.String())
	}
	if strings.Contains(rec.Body.String(), "@@BASE@@") {
		t.Fatal("placeholder left behind")
	}
}

func TestPlainHttpIsUsedWhenThereIsNoTLSOrProxyHeader(t *testing.T) {
	mux := http.NewServeMux()
	registerInstallers(mux, testFiles(), "")
	rec := get(t, mux, "/setup.sh", func(r *http.Request) { r.Host = "127.0.0.1:8080" })
	if !strings.Contains(rec.Body.String(), "http://127.0.0.1:8080/steps/ollama.sh") {
		t.Fatalf("%q", rec.Body.String())
	}
}

func TestConfiguredPublicURLWins(t *testing.T) {
	mux := http.NewServeMux()
	registerInstallers(mux, testFiles(), "https://public.example.org/")
	rec := get(t, mux, "/get", func(r *http.Request) { r.Host = "internal:8080" })
	if !strings.Contains(rec.Body.String(), "base=https://public.example.org and") {
		t.Fatalf("%q", rec.Body.String())
	}
}

func TestStepsAreServedByNameOnly(t *testing.T) {
	mux := http.NewServeMux()
	registerInstallers(mux, testFiles(), "https://x.org")
	if rec := get(t, mux, "/steps/ollama.sh", nil); rec.Code != 200 || !strings.Contains(rec.Body.String(), "ollama https://x.org") {
		t.Fatalf("%d %q", rec.Code, rec.Body.String())
	}
	for _, bad := range []string{"/steps/../setup.sh", "/steps/..%2fsetup.sh", "/steps/nope.sh", "/steps/Ollama.sh", "/steps/a.txt", "/steps/"} {
		if rec := get(t, mux, bad, nil); rec.Code == 200 {
			t.Errorf("%s must not be served", bad)
		}
	}
}

func TestLegacyAgentInstallerPathStillWorks(t *testing.T) {
	mux := http.NewServeMux()
	registerInstallers(mux, testFiles(), "https://x.org")
	if rec := get(t, mux, "/install.sh", nil); rec.Code != 200 || !strings.Contains(rec.Body.String(), "agent") {
		t.Fatalf("%d", rec.Code)
	}
	if ct := get(t, mux, "/uninstall.sh", nil).Header().Get("Content-Type"); !strings.HasPrefix(ct, "text/x-shellscript") {
		t.Fatalf("content type %q", ct)
	}
}

func TestAgentCmdNeedsLoginAndIsPerMachine(t *testing.T) {
	s := &Server{
		cfg:      Config{AgentTokens: map[string]string{"pgx": "tok-pgx", "other": "tok-other"}, PublicURL: "https://control.example.org"},
		sessions: testSession("sid"),
	}
	mux := http.NewServeMux()
	mux.HandleFunc("/api/agent-cmd", s.agentCmd)
	login := func(r *http.Request) { r.AddCookie(&http.Cookie{Name: "wxch_session", Value: "sid"}) }

	if rec := get(t, mux, "/api/agent-cmd?name=pgx", nil); rec.Code != 401 || strings.Contains(rec.Body.String(), "tok-") {
		t.Fatalf("anonymous: %d %q", rec.Code, rec.Body.String())
	}
	rec := get(t, mux, "/api/agent-cmd?name=pgx", login)
	want := "curl -fsSL https://control.example.org/get | WXCH_TOKEN=tok-pgx sh"
	if rec.Code != 200 || rec.Body.String() != want || rec.Header().Get("Cache-Control") != "no-store" {
		t.Fatalf("%d %q", rec.Code, rec.Body.String())
	}
	if rec := get(t, mux, "/api/agent-cmd?name=nope", login); rec.Code != 404 {
		t.Fatalf("unknown machine: %d", rec.Code)
	}
}

func TestRegisteredMachinesAreListedByNameOnlyAndNeedLogin(t *testing.T) {
	s := &Server{
		cfg:      Config{AgentTokens: map[string]string{"pgx": "tok-pgx", "lab": "tok-lab"}},
		sessions: testSession("sid"),
	}
	mux := http.NewServeMux()
	mux.HandleFunc("/api/agents-registered", s.agentsRegistered)
	login := func(r *http.Request) { r.AddCookie(&http.Cookie{Name: "wxch_session", Value: "sid"}) }

	if rec := get(t, mux, "/api/agents-registered", nil); rec.Code != 401 || strings.Contains(rec.Body.String(), "pgx") {
		t.Fatalf("anonymous: %d %q", rec.Code, rec.Body.String())
	}
	rec := get(t, mux, "/api/agents-registered", login)
	if rec.Code != 200 || strings.TrimSpace(rec.Body.String()) != `["lab","pgx"]` {
		t.Fatalf("%d %q", rec.Code, rec.Body.String())
	}
	if strings.Contains(rec.Body.String(), "tok-") {
		t.Fatal("a token must never be listed")
	}
}

func TestAddDiskScriptIsServedWithTheCallersAddress(t *testing.T) {
	mux := http.NewServeMux()
	registerInstallers(mux, testFiles(), "")
	rec := get(t, mux, "/add-disk.sh", func(r *http.Request) { r.Host = "control.example.org"; r.Header.Set("X-Forwarded-Proto", "https") })
	if rec.Code != 200 || !strings.Contains(rec.Body.String(), "disk https://control.example.org") {
		t.Fatalf("%d %q", rec.Code, rec.Body.String())
	}
}

func TestLabScriptsAreServedWithTheCallersAddressAndOnlyByGoodNames(t *testing.T) {
	files := testFiles()
	files["lab/comfyui-step0.sh"] = &fstest.MapFile{Data: []byte("curl @@BASE@@/x\n")}
	mux := http.NewServeMux()
	registerInstallers(mux, files, "")
	rec := get(t, mux, "/lab/comfyui-step0.sh", func(r *http.Request) { r.Host = "control.example.org"; r.Header.Set("X-Forwarded-Proto", "https") })
	if rec.Code != 200 || !strings.Contains(rec.Body.String(), "https://control.example.org/x") {
		t.Fatalf("%d %q", rec.Code, rec.Body.String())
	}
	for _, bad := range []string{"/lab/../get.sh", "/lab/nope.sh", "/lab/Comfy.sh", "/lab/a.txt"} {
		if rec := get(t, mux, bad, nil); rec.Code == 200 {
			t.Fatalf("%s must not be served", bad)
		}
	}
}

func TestHostsMapToTheirSites(t *testing.T) {
	for host, want := range map[string]string{
		"notepad.wxch.org": "notepad", "note.wxch.org": "note", "NOTE.wxch.org:8080": "note",
		"api.wxch.org": "api", "chat.wxch.org": "chat", "studio.wxch.org": "studio",
		"control.wxch.org": "", "exchange.wxch.org": "", "wxch.org": "", "www.wxch.org": "", "localhost:8080": "", "": "",
		"notes.wxch.org": "", "note": "note",
	} {
		if got := appForHost(host); got != want {
			t.Errorf("%q: got %q, want %q", host, got, want)
		}
	}
}