package utils

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"io"
	"path/filepath"
	"strings"
	"testing"
)

func createTestTarGz(t *testing.T, files map[string]string, modifyHeader func(hdr *tar.Header)) io.Reader {
	t.Helper()
	var buf bytes.Buffer
	gzw := gzip.NewWriter(&buf)
	tw := tar.NewWriter(gzw)

	for name, content := range files {
		hdr := &tar.Header{
			Name: name,
			Mode: 0644,
			Size: int64(len(content)),
		}
		if modifyHeader != nil {
			modifyHeader(hdr)
		}
		if err := tw.WriteHeader(hdr); err != nil {
			t.Fatalf("Failed to write tar header for %s: %v", name, err)
		}
		if _, err := tw.Write([]byte(content)); err != nil {
			t.Fatalf("Failed to write tar content for %s: %v", name, err)
		}
	}

	if err := tw.Close(); err != nil {
		t.Fatalf("Failed to close tar writer: %v", err)
	}
	if err := gzw.Close(); err != nil {
		t.Fatalf("Failed to close gzip writer: %v", err)
	}
	return &buf
}

func TestUntarQuadlets_Valid(t *testing.T) {
	inputFiles := map[string]string{
		"workload.kat": "kind: Workload",
		"vlb.kat":      "kind: VirtualLoadBalancer",
	}
	reader := createTestTarGz(t, inputFiles, nil)

	outputFiles, err := UntarQuadlets(reader)
	if err != nil {
		t.Fatalf("UntarQuadlets() error = %v, wantErr %v", err, false)
	}

	if len(outputFiles) != len(inputFiles) {
		t.Errorf("Expected %d files, got %d", len(inputFiles), len(outputFiles))
	}
	for name, content := range inputFiles {
		outContent, ok := outputFiles[name]
		if !ok {
			t.Errorf("Expected file %s not found in output", name)
		}
		if string(outContent) != content {
			t.Errorf("Content mismatch for %s: got '%s', want '%s'", name, string(outContent), content)
		}
	}
}

func TestUntarQuadlets_EmptyArchive(t *testing.T) {
	reader := createTestTarGz(t, map[string]string{}, nil)
	_, err := UntarQuadlets(reader)
	if err == nil {
		t.Fatal("UntarQuadlets() with empty archive did not return an error")
	}
	if !strings.Contains(err.Error(), "no .kat files found") {
		t.Errorf("Expected 'no .kat files found' error, got: %v", err)
	}
}

func TestUntarQuadlets_NonKatFile(t *testing.T) {
	inputFiles := map[string]string{"config.txt": "some data"}
	reader := createTestTarGz(t, inputFiles, nil)
	_, err := UntarQuadlets(reader)
	if err == nil {
		t.Fatal("UntarQuadlets() with non-.kat file did not return an error")
	}
	if !strings.Contains(err.Error(), "only .kat files are allowed") {
		t.Errorf("Expected 'only .kat files are allowed' error, got: %v", err)
	}
}

func TestUntarQuadlets_FileInSubdirectory(t *testing.T) {
	inputFiles := map[string]string{"subdir/workload.kat": "kind: Workload"}
	reader := createTestTarGz(t, inputFiles, nil)
	_, err := UntarQuadlets(reader)
	if err == nil {
		t.Fatal("UntarQuadlets() with file in subdirectory did not return an error")
	}
	if !strings.Contains(err.Error(), "subdirectories are not allowed") {
		t.Errorf("Expected 'subdirectories are not allowed' error, got: %v", err)
	}
}

func TestUntarQuadlets_PathTraversal(t *testing.T) {
	inputFiles := map[string]string{"../workload.kat": "kind: Workload"}
	reader := createTestTarGz(t, inputFiles, nil)
	_, err := UntarQuadlets(reader)
	if err == nil {
		t.Fatal("UntarQuadlets() with path traversal did not return an error")
	}
	if !strings.Contains(err.Error(), "contains '..'") {
		t.Errorf("Expected 'contains ..' error, got: %v", err)
	}
}

func TestUntarQuadlets_FileTooLarge(t *testing.T) {
	largeContent := strings.Repeat("a", int(maxQuadletFileSize)+1)
	inputFiles := map[string]string{"large.kat": largeContent}
	reader := createTestTarGz(t, inputFiles, nil)
	_, err := UntarQuadlets(reader)
	if err == nil {
		t.Fatal("UntarQuadlets() with large file did not return an error")
	}
	if !strings.Contains(err.Error(), "file large.kat in tar is too large") {
		t.Errorf("Expected 'file ... too large' error, got: %v", err)
	}
}

func TestUntarQuadlets_TotalSizeTooLarge(t *testing.T) {
	numFiles := (maxTotalQuadletSize / maxQuadletFileSize) + 2
	fileSize := maxQuadletFileSize / 2

	inputFiles := make(map[string]string)
	content := strings.Repeat("a", int(fileSize))
	for i := 0; i < int(numFiles); i++ {
		inputFiles[filepath.Join(".", "file"+string(rune(i+'0'))+".kat")] = content
	}

	reader := createTestTarGz(t, inputFiles, nil)
	_, err := UntarQuadlets(reader)
	if err == nil {
		t.Fatal("UntarQuadlets() with total large size did not return an error")
	}
	if !strings.Contains(err.Error(), "total size of files in tar is too large") {
		t.Errorf("Expected 'total size ... too large' error, got: %v", err)
	}
}

func TestUntarQuadlets_TooManyFiles(t *testing.T) {
	inputFiles := make(map[string]string)
	for i := 0; i <= maxQuadletFiles; i++ {
		inputFiles[filepath.Join(".", "file"+string(rune(i+'a'))+".kat")] = "content"
	}
	reader := createTestTarGz(t, inputFiles, nil)
	_, err := UntarQuadlets(reader)
	if err == nil {
		t.Fatal("UntarQuadlets() with too many files did not return an error")
	}
	if !strings.Contains(err.Error(), "too many files in quadlet bundle") {
		t.Errorf("Expected 'too many files' error, got: %v", err)
	}
}

func TestUntarQuadlets_UnsupportedFileType(t *testing.T) {
	reader := createTestTarGz(t, map[string]string{"link.kat": ""}, func(hdr *tar.Header) {
		hdr.Typeflag = tar.TypeSymlink
		hdr.Linkname = "target.kat"
		hdr.Size = 0
	})
	_, err := UntarQuadlets(reader)
	if err == nil {
		t.Fatal("UntarQuadlets() with symlink did not return an error")
	}
	if !strings.Contains(err.Error(), "unsupported file type") {
		t.Errorf("Expected 'unsupported file type' error, got: %v", err)
	}
}

func TestUntarQuadlets_CorruptedGzip(t *testing.T) {
	corruptedInput := bytes.NewBufferString("this is not a valid gzip stream")
	_, err := UntarQuadlets(corruptedInput)
	if err == nil {
		t.Fatal("UntarQuadlets() with corrupted gzip did not return an error")
	}
	if !strings.Contains(err.Error(), "failed to create gzip reader") && !strings.Contains(err.Error(), "gzip: invalid header") {
		t.Errorf("Expected 'gzip format' or 'invalid header' error, got: %v", err)
	}
}

func TestUntarQuadlets_CorruptedTar(t *testing.T) {
	var buf bytes.Buffer
	gzw := gzip.NewWriter(&buf)
	_, _ = gzw.Write([]byte("this is not a valid tar stream but inside gzip"))
	_ = gzw.Close()

	_, err := UntarQuadlets(&buf)
	if err == nil {
		t.Fatal("UntarQuadlets() with corrupted tar did not return an error")
	}
	if !strings.Contains(err.Error(), "tar") {
		t.Errorf("Expected error related to 'tar' format, got: %v", err)
	}
}