package microvm

import (
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"
)

func TestCacheProxyFallsBackOnNotFound(t *testing.T) {
	first := httptest.NewServer(http.NotFoundHandler())
	defer first.Close()
	second := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		if req.URL.Path != "/abc.narinfo" {
			t.Fatalf("path: got %q, want /abc.narinfo", req.URL.Path)
		}
		_, _ = io.WriteString(w, "ok")
	}))
	defer second.Close()

	upstreams, err := parseCacheUpstreams([]string{first.URL, second.URL})
	if err != nil {
		t.Fatal(err)
	}

	req := httptest.NewRequest(http.MethodGet, "http://guest/abc.narinfo", nil)
	rec := httptest.NewRecorder()
	cacheProxyHandler(mergeCacheUpstreams(upstreams, nil, nil), slog.Default()).ServeHTTP(rec, req)

	if rec.Code != http.StatusOK {
		t.Fatalf("status: got %d, want 200; body=%q", rec.Code, rec.Body.String())
	}
	if got := rec.Body.String(); got != "ok" {
		t.Fatalf("body: got %q, want ok", got)
	}
}

func TestCacheProxyServesNixCacheInfoItself(t *testing.T) {
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		t.Errorf("upstream should not be hit, got request for %q", req.URL.Path)
	}))
	defer upstream.Close()

	upstreams, err := parseCacheUpstreams([]string{upstream.URL})
	if err != nil {
		t.Fatal(err)
	}

	req := httptest.NewRequest(http.MethodGet, "http://guest/nix-cache-info", nil)
	rec := httptest.NewRecorder()
	cacheProxyHandler(mergeCacheUpstreams(upstreams, nil, nil), slog.Default()).ServeHTTP(rec, req)

	if rec.Code != http.StatusOK {
		t.Fatalf("status: got %d, want 200; body=%q", rec.Code, rec.Body.String())
	}
	if got := rec.Body.String(); got != nixCacheInfo {
		t.Fatalf("body: got %q, want %q", got, nixCacheInfo)
	}
}

func TestCacheProxyErrorStatusDoesNotWinRace(t *testing.T) {
	erroring := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		http.Error(w, "misdirected", http.StatusMisdirectedRequest)
	}))
	defer erroring.Close()
	healthy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		time.Sleep(50 * time.Millisecond) // lose the race to the erroring upstream
		_, _ = io.WriteString(w, "ok")
	}))
	defer healthy.Close()

	upstreams, err := parseCacheUpstreams([]string{erroring.URL, healthy.URL})
	if err != nil {
		t.Fatal(err)
	}

	req := httptest.NewRequest(http.MethodGet, "http://guest/abc.narinfo", nil)
	rec := httptest.NewRecorder()
	cacheProxyHandler(mergeCacheUpstreams(upstreams, nil, nil), slog.Default()).ServeHTTP(rec, req)

	if rec.Code != http.StatusOK {
		t.Fatalf("status: got %d, want 200; body=%q", rec.Code, rec.Body.String())
	}
	if got := rec.Body.String(); got != "ok" {
		t.Fatalf("body: got %q, want ok", got)
	}
}

func TestCacheProxyJoinsSubpathQueryAndAuth(t *testing.T) {
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		if req.URL.Path != "/sub/cache/abc.narinfo" {
			t.Errorf("path: got %q, want /sub/cache/abc.narinfo", req.URL.Path)
		}
		if got := req.URL.Query().Get("token"); got != "s3cret" {
			t.Errorf("token: got %q, want s3cret", got)
		}
		if user, pass, ok := req.BasicAuth(); !ok || user != "dawn" || pass != "woof" {
			t.Errorf("basic auth: got %q/%q/%v, want dawn/woof/true", user, pass, ok)
		}
		_, _ = io.WriteString(w, "ok")
	}))
	defer upstream.Close()

	upstreamURL := "http://dawn:woof@" + strings.TrimPrefix(upstream.URL, "http://") + "/sub/cache/?token=s3cret"
	upstreams, err := parseCacheUpstreams([]string{upstreamURL})
	if err != nil {
		t.Fatal(err)
	}

	req := httptest.NewRequest(http.MethodGet, "http://guest/abc.narinfo", nil)
	rec := httptest.NewRecorder()
	cacheProxyHandler(mergeCacheUpstreams(upstreams, nil, nil), slog.Default()).ServeHTTP(rec, req)

	if rec.Code != http.StatusOK {
		t.Fatalf("status: got %d, want 200; body=%q", rec.Code, rec.Body.String())
	}
	if got := rec.Body.String(); got != "ok" {
		t.Fatalf("body: got %q, want ok", got)
	}
}

func TestCacheProxyGuardedUpstreamCannotReachBlockedRanges(t *testing.T) {
	// httptest listens on 127.0.0.1, which is in the blocked ranges; reaching
	// it would mean a workflow-defined cache can hit the host's loopback
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		t.Errorf("guarded upstream on loopback should not be reachable, got request for %q", req.URL.Path)
	}))
	defer upstream.Close()

	upstreams, err := parseCacheUpstreams([]string{upstream.URL})
	if err != nil {
		t.Fatal(err)
	}

	req := httptest.NewRequest(http.MethodGet, "http://guest/abc.narinfo", nil)
	rec := httptest.NewRecorder()
	cacheProxyHandler(mergeCacheUpstreams(nil, upstreams, nil), slog.Default()).ServeHTTP(rec, req)

	if rec.Code != http.StatusBadGateway {
		t.Fatalf("status: got %d, want 502; body=%q", rec.Code, rec.Body.String())
	}
}

func TestCacheProxyRewritesHostHeader(t *testing.T) {
	var upstreamHost string
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		if req.Host != upstreamHost {
			t.Errorf("host: got %q, want %q", req.Host, upstreamHost)
		}
		_, _ = io.WriteString(w, "ok")
	}))
	defer upstream.Close()
	upstreamHost = strings.TrimPrefix(upstream.URL, "http://")

	upstreams, err := parseCacheUpstreams([]string{upstream.URL})
	if err != nil {
		t.Fatal(err)
	}

	req := httptest.NewRequest(http.MethodGet, "http://127.0.0.1:10500/abc.narinfo", nil)
	req.Host = "127.0.0.1:10500"
	rec := httptest.NewRecorder()
	cacheProxyHandler(mergeCacheUpstreams(upstreams, nil, nil), slog.Default()).ServeHTTP(rec, req)

	if rec.Code != http.StatusOK {
		t.Fatalf("status: got %d, want 200; body=%q", rec.Code, rec.Body.String())
	}
}

func TestReadCacheProxyInjectsExtraHeadersForTrustedOnly(t *testing.T) {
	var trustedGotAuth, guardedGotAuth string

	trustedSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		trustedGotAuth = req.Header.Get("Authorization")
		_, _ = io.WriteString(w, "from-trusted")
	}))
	defer trustedSrv.Close()

	// verify guarded upstreams never receive the injected headers
	guardedSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		guardedGotAuth = req.Header.Get("Authorization")
		// return 404 so trusted wins the race
		w.WriteHeader(http.StatusNotFound)
	}))
	defer guardedSrv.Close()

	trustedURLs, err := parseCacheUpstreams([]string{trustedSrv.URL})
	if err != nil {
		t.Fatal(err)
	}
	guardedURLs, err := parseCacheUpstreams([]string{guardedSrv.URL})
	if err != nil {
		t.Fatal(err)
	}

	extraHeaders := map[string]string{"Authorization": "Bearer readtoken"}
	handler := cacheProxyHandler(mergeCacheUpstreams(trustedURLs, guardedURLs, extraHeaders), slog.Default())

	req := httptest.NewRequest(http.MethodGet, "http://guest/abc.narinfo", nil)
	rec := httptest.NewRecorder()
	handler.ServeHTTP(rec, req)

	if rec.Code != http.StatusOK {
		t.Fatalf("status: got %d, want 200; body=%q", rec.Code, rec.Body.String())
	}
	if trustedGotAuth != "Bearer readtoken" {
		t.Fatalf("trusted Authorization: got %q, want %q", trustedGotAuth, "Bearer readtoken")
	}
	if guardedGotAuth != "" {
		t.Fatalf("guarded Authorization should be empty, got %q", guardedGotAuth)
	}
}