package zipext

import (
	"archive/tar"
	"archive/zip"
	"bufio"
	"bytes"
	"compress/flate"
	"compress/gzip"
	"errors"
	"time"
)

var errAlreadyClosed = errors.New("already closed")
var errZipNotEnabled = errors.New("zip not enabled")
var errTgzNotEnabled = errors.New("tgz not enabled")

type MemoryZip struct {
	zipEnabled bool
	zipbuffer  *bytes.Buffer
	zipwriter  *zip.Writer

	tarEnabled bool
	tarbuffer  *bytes.Buffer
	tarwriter  *tar.Writer

	open bool
}

func NewMemoryZip(enableGZip, enableTarGZ bool) *MemoryZip {

	var bz *bytes.Buffer = nil
	var z *zip.Writer = nil

	var bt *bytes.Buffer = nil
	var t *tar.Writer = nil

	if enableGZip {
		bz = new(bytes.Buffer)
		z = zip.NewWriter(bz)
	}

	if enableTarGZ {
		bt = new(bytes.Buffer)
		t = tar.NewWriter(bt)
	}

	return &MemoryZip{

		zipEnabled: enableGZip,
		zipbuffer:  bz,
		zipwriter:  z,

		tarEnabled: enableTarGZ,
		tarbuffer:  bz,
		tarwriter:  t,

		open: true,
	}
}

func (z *MemoryZip) AddFile(path string, data []byte) error {
	var err error

	if !z.open {
		return errAlreadyClosed
	}

	if z.zipEnabled {
		zipheader, err := z.zipwriter.CreateHeader(&zip.FileHeader{
			Name:     path,
			Method:   zip.Deflate,
			Modified: time.Now(),
		})
		if err != nil {
			return err
		}

		_, err = zipheader.Write(data)
		if err != nil {
			return err
		}
	}

	if z.tarEnabled {
		tarheader := &tar.Header{
			Name:     path,
			ModTime:  time.Now(),
			Typeflag: tar.TypeReg,
			Size:     int64(len(data)),
		}

		err = z.tarwriter.WriteHeader(tarheader)
		if err != nil {
			return err
		}

		_, err = z.tarwriter.Write(data)
		if err != nil {
			return err
		}
	}

	return nil
}

func (z *MemoryZip) GetZip() ([]byte, error) {
	if !z.zipEnabled {
		return nil, errZipNotEnabled
	}

	if z.open {
		err := z.Close()
		if err != nil {
			return nil, err
		}
	}

	return z.zipbuffer.Bytes(), nil
}

func (z *MemoryZip) GetTarGz() ([]byte, error) {
	if !z.tarEnabled {
		return nil, errTgzNotEnabled
	}

	if z.open {
		err := z.Close()
		if err != nil {
			return nil, err
		}
	}

	b := new(bytes.Buffer)

	gf, err := gzip.NewWriterLevel(b, flate.BestCompression)
	if err != nil {
		return nil, err
	}

	fw := bufio.NewWriter(gf)
	_, err = fw.Write(z.tarbuffer.Bytes())
	if err != nil {
		return nil, err
	}

	err = fw.Flush()
	if err != nil {
		return nil, err
	}

	err = gf.Close()
	if err != nil {
		return nil, err
	}

	return b.Bytes(), nil
}

func (z *MemoryZip) Close() error {
	if !z.open {
		return nil
	}
	z.open = false

	if z.zipEnabled {
		err := z.zipwriter.Close()
		if err != nil {
			return err
		}
	}

	if z.tarEnabled {
		err := z.tarwriter.Close()
		if err != nil {
			return err
		}

	}

	return nil
}