package setup_test

import (
	"bytes"
	"image"
	"image/color"
	"image/png"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"strings"
	"testing"

	"github.com/dariuszjaniak/warp/server/internal/config"
	"github.com/dariuszjaniak/warp/server/internal/setup"
)

// newMediaServer is a server with its media in a temporary directory.
func newMediaServer(t *testing.T, tweak func(*config.Config)) *setup.Server {
	t.Helper()
	dir := t.TempDir()
	return newTestServer(t, func(cfg *config.Config) {
		cfg.Registration.Enabled = true
		cfg.ClientAPI.RateLimit.Enabled = false
		cfg.Media.BasePath = filepath.Join(dir, "media")
		if tweak != nil {
			tweak(cfg)
		}
	})
}

// send is a request with a raw body.
func send(t *testing.T, srv *setup.Server, method, path, contentType string, body []byte, token string, header http.Header) *httptest.ResponseRecorder {
	t.Helper()
	req := httptest.NewRequest(method, path, bytes.NewReader(body))
	if contentType != "" {
		req.Header.Set("Content-Type", contentType)
	}
	if token != "" {
		req.Header.Set("Authorization", "Bearer "+token)
	}
	for k, v := range header {
		req.Header[k] = v
	}
	rec := httptest.NewRecorder()
	srv.Handler().ServeHTTP(rec, req)
	return rec
}

// uploadThrough uploads content and returns the server name and media ID
// of its mxc:// URI.
func uploadThrough(t *testing.T, srv *setup.Server, token, contentType, fileName string, body []byte) (string, string) {
	t.Helper()
	path := "/_matrix/media/v3/upload"
	if fileName != "" {
		path += "?filename=" + fileName
	}
	rec := send(t, srv, http.MethodPost, path, contentType, body, token, nil)
	if rec.Code != http.StatusOK {
		t.Fatalf("upload: %d %s", rec.Code, rec.Body.String())
	}
	return splitMXC(t, decode[map[string]string](t, rec)["content_uri"])
}

func splitMXC(t *testing.T, mxc string) (string, string) {
	t.Helper()
	rest, ok := strings.CutPrefix(mxc, "mxc://")
	server, id, found := strings.Cut(rest, "/")
	if !ok || !found {
		t.Fatalf("bad mxc URI %q", mxc)
	}
	return server, id
}

func pngOf(t *testing.T, w, h int) []byte {
	t.Helper()
	img := image.NewRGBA(image.Rect(0, 0, w, h))
	for y := range h {
		for x := range w {
			img.Set(x, y, color.RGBA{uint8(x), uint8(y), 200, 255})
		}
	}
	var buf bytes.Buffer
	if err := png.Encode(&buf, img); err != nil {
		t.Fatal(err)
	}
	return buf.Bytes()
}

func TestMediaIsWiredUp(t *testing.T) {
	srv := newMediaServer(t, nil)
	token := registerThrough(t, srv, "alice")

	rec := getWithToken(t, srv, "/_matrix/client/v1/media/config", token)
	if rec.Code != http.StatusOK || !strings.Contains(rec.Body.String(), `"m.upload.size":52428800`) {
		t.Fatalf("media config: %d %s", rec.Code, rec.Body.String())
	}

	server, id := uploadThrough(t, srv, token, "text/plain", "hello.txt", []byte("hello, media"))
	if server != "warp.example" {
		t.Errorf("uploaded to %q", server)
	}

	download := "/_matrix/client/v1/media/download/" + server + "/" + id
	rec = getWithToken(t, srv, download, token)
	if rec.Code != http.StatusOK || rec.Body.String() != "hello, media" {
		t.Fatalf("download: %d %q", rec.Code, rec.Body.String())
	}
	for header, want := range map[string]string{
		"Content-Type":           "text/plain",
		"Content-Disposition":    "inline; filename=hello.txt",
		"X-Content-Type-Options": "nosniff",
		"Cache-Control":          "private, max-age=31536000, immutable",
	} {
		if got := rec.Header().Get(header); got != want {
			t.Errorf("%s = %q, want %q", header, got, want)
		}
	}
	if csp := rec.Header().Get("Content-Security-Policy"); !strings.HasPrefix(csp, "sandbox;") {
		t.Errorf("Content-Security-Policy = %q", csp)
	}

	// The authenticated endpoints need a token.
	if rec := get(t, srv, http.MethodGet, download); rec.Code != http.StatusUnauthorized {
		t.Errorf("download without a token: %d", rec.Code)
	}
	// The deprecated unauthenticated endpoint is off by default.
	if rec := get(t, srv, http.MethodGet, "/_matrix/media/v3/download/"+server+"/"+id); rec.Code != http.StatusNotFound ||
		!strings.Contains(rec.Body.String(), "M_UNRECOGNIZED") {
		t.Errorf("legacy download: %d %s", rec.Code, rec.Body.String())
	}

	// A file name in the path renames the download.
	rec = getWithToken(t, srv, download+"/other.txt", token)
	if got := rec.Header().Get("Content-Disposition"); got != "inline; filename=other.txt" {
		t.Errorf("renamed download: Content-Disposition = %q", got)
	}

	// Ranges and conditional requests come with http.ServeContent.
	rec = send(t, srv, http.MethodGet, download, "", nil, token, http.Header{"Range": {"bytes=0-4"}})
	if rec.Code != http.StatusPartialContent || rec.Body.String() != "hello" {
		t.Errorf("range: %d %q", rec.Code, rec.Body.String())
	}
	etag := getWithToken(t, srv, download, token).Header().Get("ETag")
	rec = send(t, srv, http.MethodGet, download, "", nil, token, http.Header{"If-None-Match": {etag}})
	if rec.Code != http.StatusNotModified {
		t.Errorf("If-None-Match: %d", rec.Code)
	}

	// An HTML upload is served as an attachment, never as a page.
	_, page := uploadThrough(t, srv, token, "text/html", "page.html", []byte("<script>alert(1)</script>"))
	rec = getWithToken(t, srv, "/_matrix/client/v1/media/download/warp.example/"+page, token)
	if got := rec.Header().Get("Content-Disposition"); !strings.HasPrefix(got, "attachment") {
		t.Errorf("HTML Content-Disposition = %q", got)
	}

	rec = getWithToken(t, srv, "/_matrix/client/v1/media/download/warp.example/doesnotexist", token)
	if rec.Code != http.StatusNotFound || !strings.Contains(rec.Body.String(), "M_NOT_FOUND") {
		t.Errorf("unknown media: %d %s", rec.Code, rec.Body.String())
	}
}

func TestMediaThumbnailsThroughTheAPI(t *testing.T) {
	srv := newMediaServer(t, nil)
	token := registerThrough(t, srv, "alice")
	server, id := uploadThrough(t, srv, token, "image/png", "", pngOf(t, 400, 300))

	thumb := "/_matrix/client/v1/media/thumbnail/" + server + "/" + id
	rec := getWithToken(t, srv, thumb+"?width=64&height=64&method=crop", token)
	if rec.Code != http.StatusOK || rec.Header().Get("Content-Type") != "image/png" {
		t.Fatalf("thumbnail: %d %s %s", rec.Code, rec.Header().Get("Content-Type"), rec.Body.String())
	}
	img, err := png.Decode(rec.Body)
	if err != nil {
		t.Fatalf("thumbnail does not decode: %v", err)
	}
	if b := img.Bounds(); b.Dx() != 96 || b.Dy() != 96 {
		t.Errorf("thumbnail is %v, want 96 × 96", b)
	}

	for query, want := range map[string]string{
		"":                               "M_MISSING_PARAM",
		"?width=0&height=10":             "M_INVALID_PARAM",
		"?width=10&height=10&method=zap": "M_INVALID_PARAM",
	} {
		rec := getWithToken(t, srv, thumb+query, token)
		if rec.Code != http.StatusBadRequest || !strings.Contains(rec.Body.String(), want) {
			t.Errorf("thumbnail%s: %d %s, want %s", query, rec.Code, rec.Body.String(), want)
		}
	}

	_, text := uploadThrough(t, srv, token, "text/plain", "", []byte("words"))
	rec = getWithToken(t, srv, "/_matrix/client/v1/media/thumbnail/warp.example/"+text+"?width=32&height=32", token)
	if rec.Code != http.StatusBadRequest {
		t.Errorf("thumbnail of text: %d %s", rec.Code, rec.Body.String())
	}
}

func TestAsyncUpload(t *testing.T) {
	srv := newMediaServer(t, nil)
	alice := registerThrough(t, srv, "alice")
	bob := registerThrough(t, srv, "bob")

	rec := post(t, srv, "/_matrix/media/v1/create", map[string]any{}, alice)
	if rec.Code != http.StatusOK {
		t.Fatalf("create: %d %s", rec.Code, rec.Body.String())
	}
	created := decode[map[string]any](t, rec)
	mxc, _ := created["content_uri"].(string)
	if _, ok := created["unused_expires_at"].(float64); !ok {
		t.Errorf("create response lacks unused_expires_at: %v", created)
	}
	server, id := splitMXC(t, mxc)

	download := "/_matrix/client/v1/media/download/" + server + "/" + id
	rec = getWithToken(t, srv, download+"?timeout_ms=0", alice)
	if rec.Code != http.StatusGatewayTimeout || !strings.Contains(rec.Body.String(), "M_NOT_YET_UPLOADED") {
		t.Errorf("download before upload: %d %s", rec.Code, rec.Body.String())
	}

	put := "/_matrix/media/v3/upload/" + server + "/" + id
	if rec := send(t, srv, http.MethodPut, put, "text/plain", []byte("bob's"), bob, nil); rec.Code != http.StatusForbidden {
		t.Errorf("upload to another user's ID: %d %s", rec.Code, rec.Body.String())
	}
	if rec := send(t, srv, http.MethodPut, put, "text/plain", []byte("later"), alice, nil); rec.Code != http.StatusOK {
		t.Fatalf("upload: %d %s", rec.Code, rec.Body.String())
	}
	rec = send(t, srv, http.MethodPut, put, "text/plain", []byte("again"), alice, nil)
	if rec.Code != http.StatusConflict || !strings.Contains(rec.Body.String(), "M_CANNOT_OVERWRITE_MEDIA") {
		t.Errorf("second upload: %d %s", rec.Code, rec.Body.String())
	}
	if rec := getWithToken(t, srv, download, bob); rec.Code != http.StatusOK || rec.Body.String() != "later" {
		t.Errorf("download: %d %q", rec.Code, rec.Body.String())
	}
}

func TestUploadLimit(t *testing.T) {
	srv := newMediaServer(t, func(cfg *config.Config) { cfg.Media.MaxUploadBytes = 8 })
	token := registerThrough(t, srv, "alice")

	rec := send(t, srv, http.MethodPost, "/_matrix/media/v3/upload", "text/plain", []byte("more than eight"), token, nil)
	if rec.Code != http.StatusRequestEntityTooLarge || !strings.Contains(rec.Body.String(), "M_TOO_LARGE") {
		t.Errorf("oversized upload: %d %s", rec.Code, rec.Body.String())
	}

	// Without a Content-Length the blob store cuts it off instead.
	req := httptest.NewRequest(http.MethodPost, "/_matrix/media/v3/upload", strings.NewReader("more than eight"))
	req.ContentLength = -1
	req.Header.Set("Authorization", "Bearer "+token)
	rec = httptest.NewRecorder()
	srv.Handler().ServeHTTP(rec, req)
	if rec.Code != http.StatusRequestEntityTooLarge {
		t.Errorf("oversized chunked upload: %d %s", rec.Code, rec.Body.String())
	}
}

func TestLegacyMediaDownloads(t *testing.T) {
	srv := newMediaServer(t, func(cfg *config.Config) { cfg.Media.LegacyUnauthenticatedDownloads = true })
	token := registerThrough(t, srv, "alice")
	server, id := uploadThrough(t, srv, token, "image/png", "", pngOf(t, 10, 10))

	if rec := get(t, srv, http.MethodGet, "/_matrix/media/v3/download/"+server+"/"+id); rec.Code != http.StatusOK {
		t.Errorf("legacy download: %d %s", rec.Code, rec.Body.String())
	}
	if rec := get(t, srv, http.MethodGet, "/_matrix/media/v3/thumbnail/"+server+"/"+id+"?width=32&height=32"); rec.Code != http.StatusOK {
		t.Errorf("legacy thumbnail: %d %s", rec.Code, rec.Body.String())
	}
}

func TestMediaOverFederation(t *testing.T) {
	if testing.Short() {
		t.Skip("starts two servers")
	}
	nodes := newFederation(t, "hs-a.example", "hs-b.example")
	a, b := nodes["hs-a.example"].srv, nodes["hs-b.example"].srv
	alice := registerThrough(t, a, "alice")
	bob := registerThrough(t, b, "bob")

	server, id := uploadThrough(t, a, alice, "image/png", "cat.png", pngOf(t, 300, 200))

	// Bob's server fetches Alice's upload from her server.
	download := "/_matrix/client/v1/media/download/" + server + "/" + id
	rec := getWithToken(t, b, download, bob)
	if rec.Code != http.StatusOK || rec.Header().Get("Content-Type") != "image/png" {
		t.Fatalf("remote download: %d %s", rec.Code, rec.Body.String())
	}
	if got := rec.Header().Get("Content-Disposition"); got != "inline; filename=cat.png" {
		t.Errorf("remote download kept no file name: %q", got)
	}
	if _, err := png.Decode(rec.Body); err != nil {
		t.Errorf("remote download is not the PNG: %v", err)
	}

	// A thumbnail of media that is cached is made locally; of media that
	// is not, it is fetched.
	_, other := uploadThrough(t, a, alice, "image/png", "", pngOf(t, 500, 500))
	rec = getWithToken(t, b, "/_matrix/client/v1/media/thumbnail/"+server+"/"+other+"?width=32&height=32&method=crop", bob)
	if rec.Code != http.StatusOK {
		t.Fatalf("remote thumbnail: %d %s", rec.Code, rec.Body.String())
	}
	if img, err := png.Decode(rec.Body); err != nil || img.Bounds().Dx() != 32 {
		t.Errorf("remote thumbnail: %v, %v", img, err)
	}

	// Unknown media on the other server is a 404 here too.
	rec = getWithToken(t, b, "/_matrix/client/v1/media/download/"+server+"/doesnotexist", bob)
	if rec.Code != http.StatusNotFound {
		t.Errorf("unknown remote media: %d %s", rec.Code, rec.Body.String())
	}

	// With B's federation off, what B has cached is still served and
	// nothing else is fetched.
	if err := b.Settings().SetFederationEnabled(t.Context(), false); err != nil {
		t.Fatal(err)
	}
	if rec := getWithToken(t, b, download, bob); rec.Code != http.StatusOK {
		t.Errorf("cached remote download with federation off: %d %s", rec.Code, rec.Body.String())
	}
	_, uncached := uploadThrough(t, a, alice, "text/plain", "", []byte("never fetched"))
	rec = getWithToken(t, b, "/_matrix/client/v1/media/download/"+server+"/"+uncached, bob)
	if rec.Code != http.StatusForbidden || !strings.Contains(rec.Body.String(), "Federation is disabled") {
		t.Errorf("uncached remote download with federation off: %d %s", rec.Code, rec.Body.String())
	}

	// Another server is only served this server's own media, and only
	// over federation.
	if rec := get(t, a, http.MethodGet, "/_matrix/federation/v1/media/download/"+id); rec.Code != http.StatusUnauthorized {
		t.Errorf("unsigned federation download: %d %s", rec.Code, rec.Body.String())
	}
}