package api

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"context"
	"errors"
	"io"
	"mime"
	"net/http"
	"net/http/httptest"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"syscall"
	"testing"
)

type snapshotTransferWriter struct {
	*httptest.ResponseRecorder
	check func()
}

func (w *snapshotTransferWriter) WriteHeader(code int) {
	w.check()
	w.ResponseRecorder.WriteHeader(code)
}

func TestSnapshotDownloadTransferUnlockedAndOversize(t *testing.T) {
	a, i, home := configTestAPI(t, false)
	tmp := t.TempDir()
	t.Setenv("TMPDIR", tmp)
	id := "failed-home-" + strings.Repeat("c", 32)
	dir := filepath.Join(filepath.Dir(home), id)
	if err := os.Mkdir(dir, 0700); err != nil {
		t.Fatal(err)
	}
	configWrite(t, filepath.Join(dir, "data"), "saved")
	checked := false
	w := &snapshotTransferWriter{httptest.NewRecorder(), func() {
		checked = true
		if a.instanceLocks.held() {
			t.Fatal("instance locked during transfer")
		}
		files, err := os.ReadDir(tmp)
		if err != nil || len(files) != 1 {
			t.Fatal(files, err)
		}
		st, err := files[0].Info()
		if err != nil || st.Mode().Perm() != 0600 {
			t.Fatal(st, err)
		}
		if err := os.RemoveAll(dir); err != nil {
			t.Fatal(err)
		}
	}}
	a.snapshotDownload(w, httptest.NewRequest(http.MethodGet, "/?download="+id, nil), i)
	if !checked || w.Code != 200 {
		t.Fatal(checked, w.Code)
	}
	gz, err := gzip.NewReader(bytes.NewReader(w.Body.Bytes()))
	if err != nil {
		t.Fatal(err)
	}
	defer gz.Close()
	if _, err := io.Copy(io.Discard, gz); err != nil {
		t.Fatal(err)
	}
	if err := os.Mkdir(dir, 0700); err != nil {
		t.Fatal(err)
	}
	f, err := os.Create(filepath.Join(dir, "huge"))
	if err != nil {
		t.Fatal(err)
	}
	err = f.Truncate(upgradeSnapshotLimit + 1)
	f.Close()
	if err != nil {
		t.Fatal(err)
	}
	rec := httptest.NewRecorder()
	a.snapshotDownload(rec, httptest.NewRequest("GET", "/?download="+id, nil), i)
	if rec.Code != 413 || !strings.Contains(rec.Body.String(), "limit") {
		t.Fatal(rec.Code, rec.Body.String())
	}
	files, err := os.ReadDir(tmp)
	if err != nil || len(files) != 0 {
		t.Fatal(files, err)
	}
}

func TestSnapshotDownloadContentsAndCleanup(t *testing.T) {
	a, i, home := configTestAPI(t, true)
	tmp := t.TempDir()
	t.Setenv("TMPDIR", tmp)
	id := "pre-restore-" + strings.Repeat("a", 32)
	dir := filepath.Join(filepath.Dir(home), id)
	if err := os.MkdirAll(filepath.Join(dir, "node_modules/.pnpm/pkg"), 0750); err != nil {
		t.Fatal(err)
	}
	configWrite(t, filepath.Join(dir, "node_modules/.pnpm/pkg/bin"), "#!/bin/sh")
	if err := os.Chmod(filepath.Join(dir, "node_modules/.pnpm/pkg/bin"), 0751); err != nil {
		t.Fatal(err)
	}
	if err := os.Symlink(".pnpm/pkg", filepath.Join(dir, "node_modules/pkg")); err != nil {
		t.Fatal(err)
	}
	configWrite(t, filepath.Join(dir, ".credentials.yaml"), "secret")
	if _, err := a.s.DB.Exec("UPDATE instances SET status='error' WHERE id=?", i.ID); err != nil {
		t.Fatal(err)
	}
	w := httptest.NewRecorder()
	a.config(w, httptest.NewRequest("GET", "/?download="+id, nil), i, "snapshots")
	if w.Code != 200 {
		t.Fatal(w.Code, w.Body.String())
	}
	_, params, err := mime.ParseMediaType(w.Header().Get("Content-Disposition"))
	if err != nil || params["filename"] != id+".tar.gz" {
		t.Fatal(params, err)
	}
	gz, err := gzip.NewReader(bytes.NewReader(w.Body.Bytes()))
	if err != nil {
		t.Fatal(err)
	}
	defer gz.Close()
	tr := tar.NewReader(gz)
	seen := map[string]bool{}
	for {
		h, err := tr.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			t.Fatal(err)
		}
		seen[h.Name] = true
		b, err := io.ReadAll(tr)
		if err != nil {
			t.Fatal(err)
		}
		switch h.Name {
		case "node_modules/.pnpm/pkg/bin":
			if string(b) != "#!/bin/sh" || h.Mode != 0751 {
				t.Fatal(h, string(b))
			}
		case "node_modules/pkg":
			if h.Typeflag != tar.TypeSymlink || h.Linkname != ".pnpm/pkg" {
				t.Fatal(h)
			}
		case ".credentials.yaml":
			if string(b) != "secret" {
				t.Fatal(string(b))
			}
		}
	}
	if len(seen) != 7 || !seen["node_modules/pkg"] || seen["node_modules/pkg/bin"] {
		t.Fatal(seen)
	}
	files, err := os.ReadDir(tmp)
	if err != nil || len(files) != 0 {
		t.Fatal(files, err)
	}
	r := httptest.NewRequest("GET", "/?download="+id, nil)
	r.Header.Set("Range", "bytes=0-9")
	w = httptest.NewRecorder()
	a.snapshotDownload(w, r, i)
	if w.Code != 206 || w.Body.Len() != 10 {
		t.Fatal(w.Code, w.Body.Len())
	}
}

func TestSnapshotDownloadRejectsUnsafeAndForeign(t *testing.T) {
	a, i, home := configTestAPI(t, false)
	id := "pre-restore-" + strings.Repeat("b", 32)
	dir := filepath.Join(filepath.Dir(home), id)
	if err := os.Mkdir(dir, 0700); err != nil {
		t.Fatal(err)
	}
	for _, query := range []string{"", "../home", "/home", id + "/../home", "pre-restore-ABC"} {
		w := httptest.NewRecorder()
		a.snapshotDownload(w, httptest.NewRequest("GET", "/?download="+url.QueryEscape(query), nil), i)
		if w.Code != 400 {
			t.Fatal(query, w.Code)
		}
	}
	foreign := i
	foreign.UID = "someone-else"
	w := httptest.NewRecorder()
	a.snapshotDownload(w, httptest.NewRequest("GET", "/?download="+id, nil), foreign)
	if w.Code != 404 {
		t.Fatal(w.Code)
	}
	for _, target := range []string{"../home", "/workspace/example", "../../elsewhere", "missing", "link"} {
		if err := os.Symlink(target, filepath.Join(dir, "link")); err != nil {
			t.Fatal(err)
		}
		w = httptest.NewRecorder()
		a.snapshotDownload(w, httptest.NewRequest("GET", "/?download="+id, nil), i)
		if w.Code != 200 || w.Header().Get("Content-Disposition") == "" {
			t.Fatal(target, w.Code)
		}
		gz, err := gzip.NewReader(w.Body)
		if err != nil {
			t.Fatal(err)
		}
		tr := tar.NewReader(gz)
		found := false
		for {
			h, err := tr.Next()
			if err == io.EOF {
				break
			}
			if err != nil {
				t.Fatal(err)
			}
			if h.Name == "link" {
				found = h.Typeflag == tar.TypeSymlink && h.Linkname == target && h.Size == 0
			} else if h.Name != "./" {
				t.Fatalf("unexpected archived entry: %s", h.Name)
			}
		}
		gz.Close()
		if !found {
			t.Fatal("link text was not preserved")
		}
		if err := os.Remove(filepath.Join(dir, "link")); err != nil {
			t.Fatal(err)
		}
	}
	if err := syscall.Mkfifo(filepath.Join(dir, "fifo"), 0600); err != nil {
		t.Fatal(err)
	}
	w = httptest.NewRecorder()
	a.snapshotDownload(w, httptest.NewRequest("GET", "/?download="+id, nil), i)
	if w.Code != 400 {
		t.Fatal(w.Code)
	}
	if err := os.Remove(filepath.Join(dir, "fifo")); err != nil {
		t.Fatal(err)
	}
	if err := os.Remove(dir); err != nil {
		t.Fatal(err)
	}
	if err := os.Symlink(home, dir); err != nil {
		t.Fatal(err)
	}
	w = httptest.NewRecorder()
	a.snapshotDownload(w, httptest.NewRequest("GET", "/?download="+id, nil), i)
	if w.Code != 400 {
		t.Fatal(w.Code)
	}
}

func TestSnapshotArchiveLimitsAndCancellation(t *testing.T) {
	dir := t.TempDir()
	configWrite(t, filepath.Join(dir, "file"), strings.Repeat("x", 2048))
	root, err := os.OpenRoot(dir)
	if err != nil {
		t.Fatal(err)
	}
	defer root.Close()
	if err := snapshotArchive(context.Background(), root, io.Discard, 1024); !errors.Is(err, errSnapshotDownloadLimit) {
		t.Fatal(err)
	}
	if err := snapshotArchive(context.Background(), root, io.Discard, 2048); err != nil {
		t.Fatal("exact content limit should allow archive overhead", err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	if err := snapshotArchive(ctx, root, io.Discard, upgradeSnapshotLimit); !errors.Is(err, context.Canceled) {
		t.Fatal(err)
	}
	a, i, _ := configTestAPI(t, false)
	tmp := t.TempDir()
	t.Setenv("TMPDIR", tmp)
	r := httptest.NewRequest("GET", "/?download=pre-restore-"+strings.Repeat("a", 32), nil).WithContext(ctx)
	w := httptest.NewRecorder()
	a.snapshotDownload(w, r, i)
	if w.Code != 408 {
		t.Fatal(w.Code)
	}
	files, err := os.ReadDir(tmp)
	if err != nil || len(files) != 0 {
		t.Fatal(files, err)
	}
	if len(zipSlots) != 0 || a.instanceLocks.held() {
		t.Fatal("resource leak")
	}
	t.Setenv("TMPDIR", filepath.Join(tmp, "missing"))
	w = httptest.NewRecorder()
	a.snapshotDownload(w, httptest.NewRequest("GET", r.URL.String(), nil), i)
	if w.Code != 500 {
		t.Fatal(w.Code)
	}
}