package zip_utils

import (
	"archive/zip"
	"bytes"
	"io"
	"time"
)

func ZipData(fileName string, data []byte) ([]byte, error) {
	// Create zip
	buf := new(bytes.Buffer)

	// Create a new zip archive.
	zipWriter := zip.NewWriter(buf)

	header := &zip.FileHeader{
		Name:     fileName,
		Method:   zip.Deflate,
		Modified: time.Now(),
	}

	// Write data
	file, err := zipWriter.CreateHeader(header)
	if err != nil {
		return nil, err
	}
	_, err = file.Write(data)
	if err != nil {
		return nil, err
	}

	// Close zip
	err = zipWriter.Close()
	if err != nil {
		return []byte(""), err
	}

	return buf.Bytes(), nil
}

func UnzipData(data []byte) (map[string][]byte, error) {
	// Create a new zip reader.
	zipReader, err := zip.NewReader(bytes.NewReader(data), int64(len(data)))
	if err != nil {
		return nil, err
	}

	// Read all the files from zip archive
	var fileData = make(map[string][]byte)
	for _, zipFile := range zipReader.File {
		unzippedData, err := readZipFile(zipFile)
		if err != nil {
			continue
		}

		fileData[zipFile.FileHeader.Name] = unzippedData
	}

	return fileData, nil
}

func readZipFile(zipFile *zip.File) ([]byte, error) {
	zipFileData, err := zipFile.Open()
	if err != nil {
		return nil, err
	}
	defer zipFileData.Close()

	return io.ReadAll(zipFileData)
}

func AddFileToZip(fileName string, fileBytes []byte, zipWriter *zip.Writer) error {
	file, err := zipWriter.Create(fileName)
	if err != nil {
		return err
	}
	_, err = file.Write(fileBytes)
	if err != nil {
		return err
	}

	return nil
}