package workspace

import (
	"archive/tar"
	"bytes"
	"os"
	"path/filepath"
	"testing"
)

func createTarBuffer(files map[string]string) *bytes.Buffer {
	var buf bytes.Buffer
	tw := tar.NewWriter(&buf)
	for name, content := range files {
		hdr := &tar.Header{
			Name: name,
			Mode: 0644,
			Size: int64(len(content)),
		}
		tw.WriteHeader(hdr)
		tw.Write([]byte(content))
	}
	tw.Close()
	return &buf
}

func TestExtractTar(t *testing.T) {
	destDir := t.TempDir()

	buf := createTarBuffer(map[string]string{
		"hello.txt": "world",
		"sub/a.txt": "nested",
	})

	if err := ExtractTar(buf, destDir); err != nil {
		t.Fatalf("ExtractTar: %v", err)
	}

	data, err := os.ReadFile(filepath.Join(destDir, "hello.txt"))
	if err != nil {
		t.Fatalf("read hello.txt: %v", err)
	}
	if string(data) != "world" {
		t.Errorf("hello.txt = %q, want %q", data, "world")
	}

	data, err = os.ReadFile(filepath.Join(destDir, "sub", "a.txt"))
	if err != nil {
		t.Fatalf("read sub/a.txt: %v", err)
	}
	if string(data) != "nested" {
		t.Errorf("sub/a.txt = %q, want %q", data, "nested")
	}
}

func TestExtractTarStripPrefix(t *testing.T) {
	destDir := t.TempDir()

	buf := createTarBuffer(map[string]string{
		"toplevel/":        "",
		"toplevel/foo.txt": "bar",
	})

	if err := ExtractTarStripPrefix(buf, destDir); err != nil {
		t.Fatalf("ExtractTarStripPrefix: %v", err)
	}

	data, err := os.ReadFile(filepath.Join(destDir, "foo.txt"))
	if err != nil {
		t.Fatalf("read foo.txt: %v", err)
	}
	if string(data) != "bar" {
		t.Errorf("foo.txt = %q, want %q", data, "bar")
	}
}

func TestExtractTarStripPrefix_PathTraversal(t *testing.T) {
	destDir := t.TempDir()

	buf := createTarBuffer(map[string]string{
		"prefix/../../../etc/passwd": "pwned",
	})

	err := ExtractTarStripPrefix(buf, destDir)
	if err == nil {
		t.Fatal("expected error for path traversal, got nil")
	}
}

func TestExtractTar_PathTraversal(t *testing.T) {
	destDir := t.TempDir()

	buf := createTarBuffer(map[string]string{
		"../../../etc/passwd": "pwned",
	})

	err := ExtractTar(buf, destDir)
	if err == nil {
		t.Fatal("expected error for path traversal, got nil")
	}
}