package archiver

import (
	"archive/tar"
	"archive/zip"
	"compress/gzip"
	"github.com/dsnet/compress/bzip2"
	"github.com/go-git/go-billy/v5/memfs"
	"github.com/stretchr/testify/assert"
	"github.com/ulikunitz/xz"
	"os"
	"regexp"
	"strconv"
	"testing"
)

const (
	testArchiveBaseDir = "./test/archives"
)

var archiveRegex = regexp.MustCompile(`(?m)test[1|5]`)

func TestArchiveZip(t *testing.T) {
	err := os.MkdirAll(testArchiveBaseDir, os.ModePerm)
	assert.NoError(t, err)
	testGenerateZip(t)

	t.Log("== OpenArchive ==")
	archive, err := Open(Zip, testArchiveBaseDir+"/test.zip")
	assert.NoError(t, err)
	assert.Equal(t, archive.FileCount(), 10)

	t.Log("== ExtractArchiveFile ==")
	archiveFile, err := archive.GetFile("test0.txt")
	assert.NoError(t, err)

	err = archiveFile.Extract(ExtractFileOptions{
		Folder: testArchiveBaseDir + "/extracted/zip",
	})
	assert.NoError(t, err)

	err = os.RemoveAll(testArchiveBaseDir + "/extracted/zip")
	assert.NoError(t, err)

	t.Log("== ExtractArchive ==")
	err = archive.Extract(ExtractOptions{
		Folder: testArchiveBaseDir + "/extracted/zip",
	})
	assert.NoError(t, err)

	t.Log("== ExtractArchiveFilter ==")
	err = archive.Extract(ExtractOptions{
		Folder: testArchiveBaseDir + "/extracted/zip",
		Filter: archiveRegex,
	})

	assert.FileExists(t, testArchiveBaseDir+"/extracted/zip/test1.txt")
	assert.FileExists(t, testArchiveBaseDir+"/extracted/zip/test5.txt")

	err = os.RemoveAll(testArchiveBaseDir + "/extracted/zip")
	assert.NoError(t, err)

	err = archive.Close()
	assert.NoError(t, err)
}

func TestArchiveTarGzip(t *testing.T) {
	err := os.MkdirAll(testArchiveBaseDir, os.ModePerm)
	assert.NoError(t, err)
	testGenerateTarGzip(t)

	t.Log("== OpenArchive ==")
	archive, err := Open(TarGzip, testArchiveBaseDir+"/test.tar.gz")
	assert.NoError(t, err)
	assert.Equal(t, archive.FileCount(), 10)

	t.Log("== ExtractArchiveFile ==")
	archiveFile, err := archive.GetFile("test0.txt")
	assert.NoError(t, err)

	err = archiveFile.Extract(ExtractFileOptions{
		Folder: testArchiveBaseDir + "/extracted/gz",
	})
	assert.NoError(t, err)

	err = os.RemoveAll(testArchiveBaseDir + "/extracted/gz")
	assert.NoError(t, err)

	t.Log("== ExtractArchive ==")
	err = archive.Extract(ExtractOptions{
		Folder: testArchiveBaseDir + "/extracted/gz",
	})
	assert.NoError(t, err)

	err = os.RemoveAll(testArchiveBaseDir + "/extracted/gz")
	assert.NoError(t, err)

	t.Log("== ExtractArchiveFilter ==")
	err = archive.Extract(ExtractOptions{
		Folder: testArchiveBaseDir + "/extracted/gz",
		Filter: archiveRegex,
	})

	assert.FileExists(t, testArchiveBaseDir+"/extracted/gz/test1.txt")
	assert.FileExists(t, testArchiveBaseDir+"/extracted/gz/test5.txt")

	err = os.RemoveAll(testArchiveBaseDir + "/extracted/gz")
	assert.NoError(t, err)

	err = archive.Close()
	assert.NoError(t, err)
}

func TestArchiveTarBzip(t *testing.T) {
	err := os.MkdirAll(testArchiveBaseDir, os.ModePerm)
	assert.NoError(t, err)
	testGenerateTarBzip(t)

	t.Log("== OpenArchive ==")
	archive, err := Open(TarBzip, testArchiveBaseDir+"/test.tar.bz2")
	assert.NoError(t, err)
	assert.Equal(t, archive.FileCount(), 10)

	t.Log("== ExtractArchiveFile ==")
	archiveFile, err := archive.GetFile("test0.txt")
	assert.NoError(t, err)

	err = archiveFile.Extract(ExtractFileOptions{
		Folder: testArchiveBaseDir + "/extracted/bzip",
	})
	assert.NoError(t, err)

	err = os.RemoveAll(testArchiveBaseDir + "/extracted/bzip")
	assert.NoError(t, err)

	t.Log("== ExtractArchive ==")
	err = archive.Extract(ExtractOptions{
		Folder: testArchiveBaseDir + "/extracted/bzip",
	})
	assert.NoError(t, err)

	err = os.RemoveAll(testArchiveBaseDir + "/extracted/bzip")
	assert.NoError(t, err)

	t.Log("== ExtractArchiveFilter ==")
	err = archive.Extract(ExtractOptions{
		Folder: testArchiveBaseDir + "/extracted/bzip",
		Filter: archiveRegex,
	})

	assert.FileExists(t, testArchiveBaseDir+"/extracted/bzip/test1.txt")
	assert.FileExists(t, testArchiveBaseDir+"/extracted/bzip/test5.txt")

	err = os.RemoveAll(testArchiveBaseDir + "/extracted/bzip")
	assert.NoError(t, err)

	err = archive.Close()
	assert.NoError(t, err)
}

func TestArchiveTarXz(t *testing.T) {
	err := os.MkdirAll(testArchiveBaseDir, os.ModePerm)
	assert.NoError(t, err)
	testGenerateTarXz(t)

	t.Log("== OpenArchive ==")
	archive, err := Open(TarXz, testArchiveBaseDir+"/test.tar.xz")
	assert.NoError(t, err)
	assert.Equal(t, archive.FileCount(), 10)

	t.Log("== ExtractArchiveFile ==")
	archiveFile, err := archive.GetFile("test0.txt")
	assert.NoError(t, err)

	err = archiveFile.Extract(ExtractFileOptions{
		Folder: testArchiveBaseDir + "/extracted/xz",
	})
	assert.NoError(t, err)

	err = os.RemoveAll(testArchiveBaseDir + "/extracted/xz")
	assert.NoError(t, err)

	t.Log("== ExtractArchive ==")
	err = archive.Extract(ExtractOptions{
		Folder: testArchiveBaseDir + "/extracted/xz",
	})
	assert.NoError(t, err)

	err = os.RemoveAll(testArchiveBaseDir + "/extracted/xz")
	assert.NoError(t, err)

	t.Log("== ExtractArchiveFilter ==")
	err = archive.Extract(ExtractOptions{
		Folder: testArchiveBaseDir + "/extracted/xz",
		Filter: archiveRegex,
	})

	assert.FileExists(t, testArchiveBaseDir+"/extracted/xz/test1.txt")
	assert.FileExists(t, testArchiveBaseDir+"/extracted/xz/test5.txt")

	err = os.RemoveAll(testArchiveBaseDir + "/extracted/xz")
	assert.NoError(t, err)

	err = archive.Close()
	assert.NoError(t, err)
}

func TestArchiveTar(t *testing.T) {
	err := os.MkdirAll(testArchiveBaseDir, os.ModePerm)
	assert.NoError(t, err)
	testGenerateTar(t)

	t.Log("== OpenArchive ==")
	archive, err := Open(Tar, testArchiveBaseDir+"/test.tar")
	assert.NoError(t, err)
	assert.Equal(t, 10, archive.FileCount())

	t.Log("== ExtractArchiveFile ==")
	archiveFile, err := archive.GetFile("test0.txt")
	assert.NoError(t, err)

	err = archiveFile.Extract(ExtractFileOptions{
		Folder: testArchiveBaseDir + "/extracted/tar",
	})
	assert.NoError(t, err)

	err = os.RemoveAll(testArchiveBaseDir + "/extracted/tar")
	assert.NoError(t, err)

	t.Log("== ExtractArchive ==")
	err = archive.Extract(ExtractOptions{
		Folder: testArchiveBaseDir + "/extracted/tar",
	})
	assert.NoError(t, err)

	err = os.RemoveAll(testArchiveBaseDir + "/extracted/tar")
	assert.NoError(t, err)

	t.Log("== ExtractArchiveFilter ==")
	err = archive.Extract(ExtractOptions{
		Folder: testArchiveBaseDir + "/extracted/tar",
		Filter: archiveRegex,
	})

	assert.FileExists(t, testArchiveBaseDir+"/extracted/tar/test1.txt")
	assert.FileExists(t, testArchiveBaseDir+"/extracted/tar/test5.txt")

	err = os.RemoveAll(testArchiveBaseDir + "/extracted/tar")
	assert.NoError(t, err)

	err = archive.Close()
	assert.NoError(t, err)
}

func TestArchiveFS(t *testing.T) {
	err := os.MkdirAll(testArchiveBaseDir, os.ModePerm)
	assert.NoError(t, err)

	memoryFS := memfs.New()

	archiveF, err := memoryFS.OpenFile("archive.tar.gz", os.O_CREATE|os.O_TRUNC|os.O_WRONLY, os.ModePerm)
	assert.NoError(t, err)
	defer archiveF.Close()

	gzipWriter := gzip.NewWriter(archiveF)
	defer gzipWriter.Close()

	tarWriter := tar.NewWriter(gzipWriter)
	defer tarWriter.Close()

	for i := 0; i < 10; i++ {
		tarHeader := &tar.Header{
			Name: "test" + strconv.Itoa(i) + ".txt",
			Size: int64(len([]byte("Hello, World! #" + strconv.Itoa(i)))),
			Mode: 0600,
		}
		err = tarWriter.WriteHeader(tarHeader)
		assert.NoError(t, err)

		_, err = tarWriter.Write([]byte("Hello, World! #" + strconv.Itoa(i)))
		assert.NoError(t, err)
	}

	err = tarWriter.Close()
	assert.NoError(t, err)
	err = gzipWriter.Close()
	assert.NoError(t, err)
	err = archiveF.Close()
	assert.NoError(t, err)

	t.Log("== OpenArchive ==")
	archive, err := OpenFS(WrapBillyFS(memoryFS), TarGzip, "archive.tar.gz")
	assert.NoError(t, err)
	assert.Equal(t, 10, archive.FileCount())

	t.Log("== ExtractArchiveFile ==")
	archiveFile, err := archive.GetFile("test0.txt")
	assert.NoError(t, err)

	err = archiveFile.Extract(ExtractFileOptions{
		Folder: testArchiveBaseDir + "/extracted/fs",
	})
	assert.NoError(t, err)

	err = os.RemoveAll(testArchiveBaseDir + "/extracted/fs")
	assert.NoError(t, err)

	t.Log("== ExtractArchive ==")
	err = archive.Extract(ExtractOptions{
		Folder: testArchiveBaseDir + "/extracted/fs",
	})
	assert.NoError(t, err)

	err = os.RemoveAll(testArchiveBaseDir + "/extracted/fs")
	assert.NoError(t, err)

	t.Log("== ExtractArchiveFilter ==")
	err = archive.Extract(ExtractOptions{
		Folder: testArchiveBaseDir + "/extracted/fs",
		Filter: archiveRegex,
	})

	assert.FileExists(t, testArchiveBaseDir+"/extracted/fs/test1.txt")
	assert.FileExists(t, testArchiveBaseDir+"/extracted/fs/test5.txt")

	err = os.RemoveAll(testArchiveBaseDir + "/extracted/fs")
	assert.NoError(t, err)

	err = archive.Close()
	assert.NoError(t, err)
}

func TestExtractArchiveBillyFS(t *testing.T) {
	err := os.MkdirAll(testArchiveBaseDir, os.ModePerm)
	assert.NoError(t, err)
	testGenerateTar(t)

	archive, err := Open(Tar, testArchiveBaseDir+"/test.tar")
	assert.NoError(t, err)

	memoryFS := memfs.New()

	err = archive.ExtractBillyFS(memoryFS, ExtractOptions{
		Folder: memoryFS.Root(),
	})
	assert.NoError(t, err)

	files, err := memoryFS.ReadDir(memoryFS.Root())
	assert.NoError(t, err)

	assert.Equal(t, 10, len(files))

	for _, file := range files {
		assert.NotEqual(t, int64(0), file.Size())
	}
}

func testGenerateTar(t *testing.T) {
	tarFile, err := os.Create(testArchiveBaseDir + "/test.tar")
	defer tarFile.Close()
	assert.NoError(t, err)

	tarWriter := tar.NewWriter(tarFile)
	defer tarWriter.Close()

	for i := 0; i < 10; i++ {
		tarHeader := &tar.Header{
			Name: "test" + strconv.Itoa(i) + ".txt",
			Size: int64(len([]byte("Hello, World! #" + strconv.Itoa(i)))),
			Mode: 0600,
		}
		err = tarWriter.WriteHeader(tarHeader)
		assert.NoError(t, err)

		_, err = tarWriter.Write([]byte("Hello, World! #" + strconv.Itoa(i)))
		assert.NoError(t, err)
	}

	err = tarWriter.Close()
	assert.NoError(t, err)
	err = tarFile.Close()
	assert.NoError(t, err)
}

func testGenerateTarXz(t *testing.T) {
	tarXzFile, err := os.Create(testArchiveBaseDir + "/test.tar.xz")
	defer tarXzFile.Close()
	assert.NoError(t, err)

	xzWriter, err := xz.NewWriter(tarXzFile)
	defer xzWriter.Close()
	assert.NoError(t, err)

	tarWriter := tar.NewWriter(xzWriter)
	defer tarWriter.Close()

	for i := 0; i < 10; i++ {
		tarHeader := &tar.Header{
			Name: "test" + strconv.Itoa(i) + ".txt",
			Size: int64(len([]byte("Hello, World! #" + strconv.Itoa(i)))),
			Mode: 0600,
		}
		err = tarWriter.WriteHeader(tarHeader)
		assert.NoError(t, err)

		_, err = tarWriter.Write([]byte("Hello, World! #" + strconv.Itoa(i)))
		assert.NoError(t, err)
	}

	err = tarWriter.Close()
	assert.NoError(t, err)
	err = xzWriter.Close()
	assert.NoError(t, err)
	err = tarXzFile.Close()
	assert.NoError(t, err)
}

func testGenerateTarBzip(t *testing.T) {
	tarBzFile, err := os.Create(testArchiveBaseDir + "/test.tar.bz2")
	defer tarBzFile.Close()
	assert.NoError(t, err)

	bzipWriter, err := bzip2.NewWriter(tarBzFile, &bzip2.WriterConfig{})
	defer bzipWriter.Close()
	assert.NoError(t, err)

	tarWriter := tar.NewWriter(bzipWriter)
	defer tarWriter.Close()

	for i := 0; i < 10; i++ {
		tarHeader := &tar.Header{
			Name: "test" + strconv.Itoa(i) + ".txt",
			Size: int64(len([]byte("Hello, World! #" + strconv.Itoa(i)))),
			Mode: 0600,
		}
		err = tarWriter.WriteHeader(tarHeader)
		assert.NoError(t, err)

		_, err = tarWriter.Write([]byte("Hello, World! #" + strconv.Itoa(i)))
		assert.NoError(t, err)
	}

	err = tarWriter.Close()
	assert.NoError(t, err)
	err = bzipWriter.Close()
	assert.NoError(t, err)
	err = tarBzFile.Close()
	assert.NoError(t, err)
}

func testGenerateTarGzip(t *testing.T) {
	tarGzFile, err := os.Create(testArchiveBaseDir + "/test.tar.gz")
	defer tarGzFile.Close()
	assert.NoError(t, err)

	gzipWriter := gzip.NewWriter(tarGzFile)
	defer gzipWriter.Close()

	tarWriter := tar.NewWriter(gzipWriter)
	defer tarWriter.Close()

	for i := 0; i < 10; i++ {
		tarHeader := &tar.Header{
			Name: "test" + strconv.Itoa(i) + ".txt",
			Size: int64(len([]byte("Hello, World! #" + strconv.Itoa(i)))),
			Mode: 0600,
		}
		err = tarWriter.WriteHeader(tarHeader)
		assert.NoError(t, err)

		_, err = tarWriter.Write([]byte("Hello, World! #" + strconv.Itoa(i)))
		assert.NoError(t, err)
	}

	err = tarWriter.Close()
	assert.NoError(t, err)
	err = gzipWriter.Close()
	assert.NoError(t, err)
	err = tarGzFile.Close()
	assert.NoError(t, err)
}

func testGenerateZip(t *testing.T) {
	zipFile, err := os.Create(testArchiveBaseDir + "/test.zip")
	defer zipFile.Close()
	assert.NoError(t, err)

	zipWriter := zip.NewWriter(zipFile)
	for i := 0; i < 10; i++ {
		file, err := zipWriter.Create("test" + strconv.Itoa(i) + ".txt")
		assert.NoError(t, err)
		_, err = file.Write([]byte("Hello, World! #" + strconv.Itoa(i)))
		assert.NoError(t, err)
	}
	err = zipWriter.Close()
	assert.NoError(t, err)
	err = zipFile.Close()
	assert.NoError(t, err)
}