package api

import (
	"archive/zip"
	"bytes"
	"context"
	"errors"
	"fmt"
	"git.misaka.ren/m1saka/dsh-web/internal/store"
	"io"
	"mime/multipart"
	"net/http"
	"net/http/httptest"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"syscall"
	"testing"
	"time"
)

func TestUploadExtensionPolicy(t *testing.T) {
	for _, tc := range []struct {
		name, policy string
		want         bool
	}{
		{"file.exe", "", true}, {"file.TXT", ".txt, .md", true},
		{"file.txt.exe", ".txt", false}, {"file", ".txt", false},
		{"file.exe", "*,txt", false}, {"file.json", ".txt,.json", true},
	} {
		if got := uploadExtensionAllowed(tc.name, tc.policy); got != tc.want {
			t.Errorf("%q %q: %v", tc.name, tc.policy, got)
		}
	}
}

func downloadTestRoot(t *testing.T) (string, *os.Root) {
	t.Helper()
	dir := t.TempDir()
	root, err := os.OpenRoot(dir)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { root.Close() })
	return dir, root
}

func TestDirectoryDownloadZip(t *testing.T) {
	dir, root := downloadTestRoot(t)
	if err := os.MkdirAll(filepath.Join(dir, "docs/empty"), 0700); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(filepath.Join(dir, "docs/你好.txt"), []byte("hello zip"), 0600); err != nil {
		t.Fatal(err)
	}
	for _, p := range []string{"docs", ".", "", "docs/empty"} {
		t.Run(p, func(t *testing.T) {
			w := httptest.NewRecorder()
			(&API{}).wsDownload(w, httptest.NewRequest("GET", "/api/ws/download?path="+url.QueryEscape(p), nil), root)
			if w.Code != 200 || w.Header().Get("Content-Type") != "application/zip" || !strings.Contains(w.Header().Get("Content-Disposition"), ".zip") {
				t.Fatal(w.Code, w.Header(), w.Body.String())
			}
			z, err := zip.NewReader(bytes.NewReader(w.Body.Bytes()), int64(w.Body.Len()))
			if err != nil {
				t.Fatal(err)
			}
			want := map[string]string{"empty/": "", "你好.txt": "hello zip"}
			if p == "." || p == "" {
				want = map[string]string{"docs/": "", "docs/empty/": "", "docs/你好.txt": "hello zip"}
			}
			if p == "docs/empty" {
				want = map[string]string{}
			}
			if len(z.File) != len(want) {
				t.Fatalf("entries %d want %d", len(z.File), len(want))
			}
			for _, f := range z.File {
				expect, ok := want[f.Name]
				if !ok {
					t.Fatal("unexpected entry", f.Name)
				}
				r, err := f.Open()
				if err != nil {
					t.Fatal(err)
				}
				b, err := io.ReadAll(r)
				r.Close()
				if err != nil || string(b) != expect {
					t.Fatal(f.Name, string(b), err)
				}
			}
		})
	}
	r := httptest.NewRequest("GET", "/api/ws/download?path=docs/"+url.QueryEscape("你好.txt"), nil)
	r.Header.Set("Range", "bytes=1-3")
	w := httptest.NewRecorder()
	(&API{}).wsDownload(w, r, root)
	if w.Code != 206 || w.Body.String() != "ell" {
		t.Fatal(w.Code, w.Body.String())
	}
}

func TestDirectoryDownloadRejectsUnsafeEntries(t *testing.T) {
	for _, kind := range []string{"symlink", "internal-link", "hardlink", "fifo", "backslash", "drive"} {
		t.Run(kind, func(t *testing.T) {
			dir, root := downloadTestRoot(t)
			if err := os.Mkdir(filepath.Join(dir, "docs"), 0700); err != nil {
				t.Fatal(err)
			}
			outside := filepath.Join(t.TempDir(), "secret")
			if err := os.WriteFile(outside, []byte("secret"), 0600); err != nil {
				t.Fatal(err)
			}
			target := filepath.Join(dir, "docs/unsafe")
			var err error
			switch kind {
			case "symlink":
				err = os.Symlink(outside, target)
			case "internal-link":
				err = os.Symlink(".", target)
			case "hardlink":
				err = os.Link(outside, target)
			case "fifo":
				err = syscall.Mkfifo(target, 0600)
			case "backslash":
				err = os.WriteFile(filepath.Join(dir, `docs/..\secret`), nil, 0600)
			case "drive":
				err = os.WriteFile(filepath.Join(dir, "docs/C:secret"), nil, 0600)
			}
			if err != nil {
				t.Fatal(err)
			}
			w := httptest.NewRecorder()
			(&API{}).wsDownload(w, httptest.NewRequest("GET", "/api/ws/download?path=docs", nil), root)
			if w.Code != 400 || strings.Contains(w.Body.String(), "secret") {
				t.Fatal(w.Code, w.Body.String())
			}
		})
	}
	dir, root := downloadTestRoot(t)
	if err := os.Mkdir(filepath.Join(dir, "real"), 0700); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(filepath.Join(dir, "real/file"), []byte("secret"), 0600); err != nil {
		t.Fatal(err)
	}
	if err := os.Symlink("real", filepath.Join(dir, "link")); err != nil {
		t.Fatal(err)
	}
	for _, p := range []string{"../secret", "/etc/passwd", `..\secret`, "link", "link/file"} {
		w := httptest.NewRecorder()
		(&API{}).wsDownload(w, httptest.NewRequest("GET", "/api/ws/download?path="+url.QueryEscape(p), nil), root)
		if w.Code < 400 {
			t.Fatal("unsafe path accepted", p)
		}
	}
}

func TestDirectoryDownloadLimits(t *testing.T) {
	for _, limit := range []string{"size", "depth", "count"} {
		t.Run(limit, func(t *testing.T) {
			dir, root := downloadTestRoot(t)
			switch limit {
			case "size":
				f, err := os.Create(filepath.Join(dir, "large"))
				if err != nil {
					t.Fatal(err)
				}
				err = f.Truncate(maxZipBytes + 1)
				f.Close()
				if err != nil {
					t.Fatal(err)
				}
			case "depth":
				if err := os.MkdirAll(filepath.Join(dir, strings.Repeat("d/", maxZipDepth+1)), 0700); err != nil {
					t.Fatal(err)
				}
			case "count":
				for i := 0; i <= maxZipEntries; i++ {
					if err := os.WriteFile(filepath.Join(dir, fmt.Sprint(i)), nil, 0600); err != nil {
						t.Fatal(err)
					}
				}
			}
			w := httptest.NewRecorder()
			(&API{}).wsDownload(w, httptest.NewRequest("GET", "/api/ws/download?path=.", nil), root)
			if w.Code != 413 {
				t.Fatal(w.Code, w.Body.String())
			}
		})
	}
}

func TestZipManifestRechecksAndCancellation(t *testing.T) {
	for _, change := range []string{"symlink", "grow", "cancel"} {
		t.Run(change, func(t *testing.T) {
			dir, root := downloadTestRoot(t)
			p := filepath.Join(dir, "file")
			if err := os.WriteFile(p, []byte("data"), 0600); err != nil {
				t.Fatal(err)
			}
			f, err := root.Open(".")
			if err != nil {
				t.Fatal(err)
			}
			defer f.Close()
			ctx, cancel := context.WithCancel(context.Background())
			defer cancel()
			es, err := zipManifest(ctx, f)
			if err != nil {
				t.Fatal(err)
			}
			switch change {
			case "symlink":
				if err := os.Remove(p); err != nil {
					t.Fatal(err)
				}
				if err := os.Symlink("/etc/passwd", p); err != nil {
					t.Fatal(err)
				}
			case "grow":
				if err := os.WriteFile(p, []byte("more data"), 0600); err != nil {
					t.Fatal(err)
				}
			case "cancel":
				cancel()
			}
			if err := streamZip(ctx, io.Discard, f, es); err == nil {
				t.Fatal("unsafe stream accepted")
			}
			if change == "cancel" {
				if _, err := zipManifest(ctx, f); !errors.Is(err, context.Canceled) {
					t.Fatal(err)
				}
			}
		})
	}
}

type failingZipWriter struct{}

func (failingZipWriter) Write([]byte) (int, error) { return 0, io.ErrClosedPipe }

func TestZipWriteFailureAndSizeBound(t *testing.T) {
	dir, root := downloadTestRoot(t)
	p := filepath.Join(dir, "file")
	if err := os.WriteFile(p, bytes.Repeat([]byte("x"), 65536), 0600); err != nil {
		t.Fatal(err)
	}
	f, err := root.Open(".")
	if err != nil {
		t.Fatal(err)
	}
	defer f.Close()
	es, err := zipManifest(context.Background(), f)
	if err != nil {
		t.Fatal(err)
	}
	if err := streamZip(context.Background(), failingZipWriter{}, f, es); !errors.Is(err, io.ErrClosedPipe) {
		t.Fatal(err)
	}
	w := &growingZipWriter{path: p}
	if err := streamZip(context.Background(), w, f, es); err == nil {
		t.Fatal("concurrent growth accepted")
	}
	if w.Len() > 65536+1024 {
		t.Fatal("stream exceeded manifest bound", w.Len())
	}
}

type growingZipWriter struct {
	bytes.Buffer
	path string
}

func (w *growingZipWriter) Write(p []byte) (int, error) {
	if w.path != "" {
		f, err := os.OpenFile(w.path, os.O_WRONLY|os.O_APPEND, 0)
		if err != nil {
			return 0, err
		}
		_, err = f.Write([]byte("growth"))
		f.Close()
		if err != nil {
			return 0, err
		}
		w.path = ""
	}
	return w.Buffer.Write(p)
}

type blockedZipWriter struct {
	*httptest.ResponseRecorder
	started   chan struct{}
	release   chan struct{}
	deadlines []time.Time
}

func (w *blockedZipWriter) Write(p []byte) (int, error) {
	if w.started != nil {
		close(w.started)
		w.started = nil
		<-w.release
	}
	return w.ResponseRecorder.Write(p)
}
func (w *blockedZipWriter) SetWriteDeadline(d time.Time) error {
	w.deadlines = append(w.deadlines, d)
	return nil
}

func TestZipTransferNoLocksAndConcurrencyLimit(t *testing.T) {
	_, root := downloadTestRoot(t)
	a := &API{}
	started := make(chan struct{})
	w := &blockedZipWriter{ResponseRecorder: httptest.NewRecorder(), started: started, release: make(chan struct{})}
	done := make(chan struct{})
	go func() {
		defer close(done)
		a.wsDownload(w, httptest.NewRequest("GET", "/api/ws/download?path=.", nil), root)
	}()
	defer func() { close(w.release); <-done }()
	select {
	case <-started:
	case <-time.After(2 * time.Second):
		t.Fatal("download not started")
	}
	if a.userLocks.held() {
		t.Fatal("user write lock held across write")
	}
	if a.instanceLocks.held() {
		t.Fatal("instance lock held across write")
	}
	if len(w.deadlines) != 1 || w.deadlines[0].IsZero() {
		t.Fatal("missing deadline")
	}
	for i := 1; i < cap(zipSlots); i++ {
		zipSlots <- struct{}{}
	}
	defer func() {
		for i := 1; i < cap(zipSlots); i++ {
			<-zipSlots
		}
	}()
	busy := httptest.NewRecorder()
	a.wsDownload(busy, httptest.NewRequest("GET", "/api/ws/download?path=.", nil), root)
	if busy.Code != 429 {
		t.Fatal(busy.Code)
	}
}

func TestRenameNoReplace(t *testing.T) {
	dir := t.TempDir()
	root, err := os.OpenRoot(dir)
	if err != nil {
		t.Fatal(err)
	}
	defer root.Close()
	for _, source := range []string{"source", ".upload-test"} {
		if err := os.WriteFile(filepath.Join(dir, source), []byte("new"), 0600); err != nil {
			t.Fatal(err)
		}
		if _, err := root.Lstat("target"); !errors.Is(err, os.ErrNotExist) {
			t.Fatal(err)
		}
		if err := os.WriteFile(filepath.Join(dir, "target"), []byte("runtime"), 0600); err != nil {
			t.Fatal(err)
		}
		if err := renameNoReplace(root, source, "target"); !errors.Is(err, os.ErrExist) {
			t.Fatalf("overwrite: %v", err)
		}
		b, err := root.ReadFile("target")
		if err != nil || string(b) != "runtime" {
			t.Fatalf("target changed: %q %v", b, err)
		}
		if err := root.Remove("target"); err != nil {
			t.Fatal(err)
		}
		if err := renameNoReplace(root, source, "target"); err != nil {
			t.Fatal(err)
		}
		if err := root.Remove("target"); err != nil {
			t.Fatal(err)
		}
	}
	if err := os.Mkdir(filepath.Join(dir, "real"), 0700); err != nil {
		t.Fatal(err)
	}
	if err := os.Symlink("real", filepath.Join(dir, "link")); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(filepath.Join(dir, "source"), []byte("new"), 0600); err != nil {
		t.Fatal(err)
	}
	for _, target := range []string{"link/target", "../target"} {
		if err := renameNoReplace(root, "source", target); err == nil {
			t.Fatalf("unsafe parent accepted: %s", target)
		}
	}
	if err := renameNoReplace(root, "link", "target"); err == nil {
		t.Fatal("symlink source accepted")
	}
}

type stalledUpload struct {
	started chan struct{}
	release chan struct{}
}

func (b *stalledUpload) Read([]byte) (int, error) {
	close(b.started)
	<-b.release
	return 0, io.EOF
}
func (b *stalledUpload) Close() error { return nil }

type deadlineRecorder struct {
	*httptest.ResponseRecorder
	deadlines []time.Time
}

func (w *deadlineRecorder) SetReadDeadline(d time.Time) error {
	w.deadlines = append(w.deadlines, d)
	return nil
}

func TestLegacyWorkspaceRemoved(t *testing.T) {
	a, sid := runtimeTestAPI(t, func(w http.ResponseWriter, r *http.Request) { t.Fatal("unexpected docker call") })
	for _, method := range []string{"GET", "POST", "DELETE"} {
		for _, endpoint := range []string{"tree", "download", "upload", "mkdir", "rename", "delete"} {
			w := httptest.NewRecorder()
			a.workspace(w, runtimeRequest(method, "/api/ws/"+endpoint, "", sid))
			if w.Code != http.StatusGone {
				t.Fatalf("%s %s: %d", method, endpoint, w.Code)
			}
		}
	}
}

func TestWorkspaceSafetyAndUpload(t *testing.T) {
	s, err := store.Open(filepath.Join(t.TempDir(), "test.sqlite"))
	if err != nil {
		t.Fatal(err)
	}
	defer s.Close()
	a := &API{s: s, root: t.TempDir()}
	u := store.User{ID: "user", Quota: 8}
	if err := s.CreateUser(u, "unused"); err != nil {
		t.Fatal(err)
	}
	if err := os.MkdirAll(filepath.Join(a.root, "workspaces/user/instances/one/home"), 0750); err != nil {
		t.Fatal(err)
	}
	root, e := a.openWorkspace(u, "one")
	if e != nil {
		t.Fatal(e)
	}
	defer root.Close()
	for _, p := range []string{"", ".", "a/..", "../other", "/tmp"} {
		w := httptest.NewRecorder()
		a.wsDelete(w, httptest.NewRequest("DELETE", "/api/ws/delete?path="+p, nil), root, u)
		if w.Code < 400 {
			t.Fatalf("deleted %q", p)
		}
	}
	outside := t.TempDir()
	os.WriteFile(filepath.Join(outside, "secret"), []byte("secret"), 0600)
	if e = os.Symlink(outside, filepath.Join(a.root, "workspaces/user/instances/one/home/link")); e != nil {
		t.Fatal(e)
	}
	w := httptest.NewRecorder()
	a.wsDownload(w, httptest.NewRequest("GET", "/api/ws/download?path=link/secret", nil), root)
	if w.Code < 400 {
		t.Fatal("escaped root")
	}
	upload := func(name, content string) int {
		var b bytes.Buffer
		m := multipart.NewWriter(&b)
		m.WriteField("path", ".")
		f, _ := m.CreateFormFile("file", name)
		f.Write([]byte(content))
		m.Close()
		r := httptest.NewRequest("POST", "/api/ws/upload", &b)
		r.Header.Set("Content-Type", m.FormDataContentType())
		w := httptest.NewRecorder()
		a.wsUpload(w, r, root, u)
		return w.Code
	}
	if code := upload("file.txt", "hello"); code != 200 {
		t.Fatal(code)
	}
	if code := upload("big.txt", "12345"); code != 413 {
		t.Fatal("quota", code)
	}
	if code := upload("file.txt", "x"); code != 409 {
		t.Fatal("overwrite", code)
	}
	w = httptest.NewRecorder()
	a.wsDownload(w, httptest.NewRequest("GET", "/api/ws/download?path=file.txt", nil), root)
	if w.Code != 200 || w.Body.String() != "hello" {
		t.Fatal(w.Code, w.Body.String())
	}
	for _, op := range []struct{ action, body string }{
		{"mkdir", `{"path":"folder"}`},
		{"rename", `{"path":"folder","new_path":"renamed"}`},
		{"delete", ""},
	} {
		w := httptest.NewRecorder()
		r := httptest.NewRequest("POST", "/api/ws/"+op.action, strings.NewReader(op.body))
		switch op.action {
		case "mkdir":
			a.wsMkdir(w, r, root, u)
		case "rename":
			a.wsRename(w, r, root, u)
		case "delete":
			a.wsDelete(w, httptest.NewRequest("DELETE", "/api/ws/delete?path=renamed", nil), root, u)
		}
		if w.Code != 200 {
			t.Fatal(op.action, w.Code, w.Body.String())
		}
	}
	rows, err := s.DB.Query("SELECT uid,action,target,detail FROM audit ORDER BY id")
	if err != nil {
		t.Fatal(err)
	}
	defer rows.Close()
	want := [][3]string{{"workspace.upload", "file.txt", ""}, {"workspace.mkdir", "folder", ""}, {"workspace.rename", "folder", "renamed"}, {"workspace.delete", "renamed", ""}}
	i := 0
	for rows.Next() {
		var uid, action, target, detail string
		if err := rows.Scan(&uid, &action, &target, &detail); err != nil {
			t.Fatal(err)
		}
		if i >= len(want) || uid != u.ID || [3]string{action, target, detail} != want[i] {
			t.Fatal("unexpected audit", uid, action, target, detail)
		}
		i++
	}
	if err := rows.Err(); err != nil {
		t.Fatal(err)
	}
	if i != len(want) {
		t.Fatal("missing audit entries", i)
	}
}