package archiver

import (
	"archive/tar"
	"archive/zip"
	"compress/gzip"
	"github.com/dsnet/compress/bzip2"
	"github.com/stretchr/testify/assert"
	"github.com/ulikunitz/xz"
	"os"
	"regexp"
	"strconv"
	"sync"
	"testing"
)

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

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

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

	testGenerateZip(t)

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

	var wg sync.WaitGroup

	t.Log("== ExtractArchiveFile ==")

	// Slam it, see if it breaks.
	for _ = range 10_000 {
		go func() {
			wg.Add(1)
			defer wg.Done()
			file, err := archive.GetFile("test0.txt")
			assert.NoError(t, err)
			assert.Equal(t, "test0.txt", file.FileName)

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

	wg.Wait()

	t.Log("== ExtractArchive ==")
	// Slam it, see if it breaks.
	for _ = range 10_000 {
		go func() {
			wg.Add(1)
			defer wg.Done()

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

	wg.Wait()

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

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

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

	testGenerateZip(t)
	testGenerateTar(t)
	testGenerateTarXz(t)
	testGenerateTarBzip(t)
	TestArchiveTarGzip(t)

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

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

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

	t.Log("== ExtractTarBz ==")
	err = Extract(testArchiveBaseDir+"/test.tar.bz2", ExtractOptions{
		Folder: testArchiveBaseDir + "/extracted/bzip",
	})
	assert.NoError(t, err)

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

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

}

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

	t.Log("== OpenArchive ==")
	archive, err := Open(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(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(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(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(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 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)
}