package archive

import (
	"archive/tar"
	"bytes"
	"errors"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/gravitational/gravity/lib/defaults"

	dockerarchive "github.com/docker/docker/pkg/archive"
	"github.com/gravitational/trace"
	log "github.com/sirupsen/logrus"
)

func DefaultOptions() *dockerarchive.TarOptions {
	options := &dockerarchive.TarOptions{NoLchown: true}
	return options
}

func CompressDirectory(dir string, writer io.Writer, items ...*Item) error {
	archive := NewTarAppender(writer)
	defer archive.Close()

	if err := archive.Add(items...); err != nil {
		return trace.Wrap(err, "failed to write tarball: %v", err.Error())
	}
	if err := filepath.Walk(dir, func(path string, fi os.FileInfo, err error) error {
		if err != nil {
			return trace.Wrap(err)
		}
		localPath, err := filepath.Rel(dir, path)
		if err != nil {
			return trace.Wrap(err)
		}
		if localPath == "." {

			return nil
		}
		item, err := ItemFromFile(localPath, path, fi)
		if err != nil {
			return trace.Wrap(err)
		}
		return trace.Wrap(archive.Add(item))
	}); err != nil {
		return trace.Wrap(err, "failed to compress directory %q", dir)
	}
	return nil
}

func Unpack(path string) (unpackedDir string, err error) {
	file, err := os.Open(path)
	if err != nil {
		return "", trace.ConvertSystemError(err)
	}
	defer file.Close()
	tmp, err := ioutil.TempDir("", "")
	if err != nil {
		return "", trace.ConvertSystemError(err)
	}
	err = Extract(file, tmp)
	if err != nil {
		return "", trace.Wrap(err)
	}
	return tmp, nil
}

func Extract(r io.Reader, dir string) error {
	tarball := tar.NewReader(r)

	for {
		header, err := tarball.Next()
		if err == io.EOF {
			break
		} else if err != nil {
			return trace.Wrap(err)
		}

		if err := extractFile(tarball, header, dir); err != nil {
			return trace.Wrap(err)
		}
	}
	return nil
}

func HasFile(tarballPath, filename string) error {
	file, err := os.Open(tarballPath)
	if err != nil {
		return trace.ConvertSystemError(err)
	}
	defer file.Close()
	var hasFile bool
	err = TarGlob(tar.NewReader(file), ".", []string{filename},
		func(match string, file io.Reader) error {
			hasFile = true
			return Abort
		})
	if err != nil {
		if trace.Unwrap(err) == tar.ErrHeader {
			return trace.BadParameter("file %v does not appear to be a valid tarball",
				tarballPath)
		}
		return trace.Wrap(err)
	}
	if !hasFile {
		return trace.NotFound("tarball %v does not contain file %v",
			tarballPath, filename)
	}
	return nil
}

func TarGlob(source *tar.Reader, dir string, patterns []string, handler func(match string, file io.Reader) error) (err error) {
	for {
		var hdr *tar.Header
		hdr, err = source.Next()
		if err == io.EOF {
			err = nil
			break
		}
		if err != nil {
			return trace.Wrap(err)
		}
		if hdr.FileInfo().IsDir() {
			continue
		}
		for _, pattern := range patterns {
			relpath, err := filepath.Rel(dir, hdr.Name)
			if err == nil {
				matched, _ := filepath.Match(pattern, filepath.Base(relpath))
				if matched {
					if err = handler(relpath, source); err != nil {
						if trace.Unwrap(err) == Abort {
							return nil
						}
						return trace.Wrap(err)
					}
				}
			}
		}
	}
	return nil
}

type TarAppender struct {
	tw *tar.Writer
}

func NewTarAppender(w io.Writer) *TarAppender {
	return &TarAppender{tar.NewWriter(w)}
}

func (r *TarAppender) Add(items ...*Item) (err error) {
	defer func() {
		for _, item := range items {
			if item.Data != nil {
				item.Data.Close()
			}
		}
	}()
	for _, item := range items {
		if item.ModTime.IsZero() {
			item.ModTime = time.Now()
		}
		if err = r.tw.WriteHeader(&item.Header); err != nil {
			return trace.Wrap(err)
		}
		if item.Data == nil {
			continue
		}
		if _, err = io.Copy(r.tw, item.Data); err != nil {
			return trace.Wrap(err)
		}
	}
	return nil
}

func (r *TarAppender) Close() error {
	return r.tw.Close()
}

func ItemFromString(path, value string) *Item {
	return ItemFromStringMode(path, value, defaults.SharedExecutableMask)
}

func DirItem(path string) *Item {
	return &Item{
		Header: tar.Header{
			Name:     path + "/",
			Typeflag: tar.TypeDir,
			Mode:     defaults.SharedDirMask,
			Uid:      defaults.ArchiveUID,
			Gid:      defaults.ArchiveGID,
		},
	}
}

func ItemFromStringMode(path, value string, mode int64) *Item {
	return ItemFromStream(path, ioutil.NopCloser(strings.NewReader(value)), int64(len(value)), mode)
}

func ItemFromStream(path string, rc io.ReadCloser, size, mode int64) *Item {
	return &Item{
		Header: tar.Header{
			Name: path,
			Size: size,
			Mode: mode,
			Uid:  defaults.ArchiveUID,
			Gid:  defaults.ArchiveGID,
		},
		Data: rc,
	}
}

func ItemFromFile(localPath, path string, fi os.FileInfo) (*Item, error) {
	fiHeader, err := tar.FileInfoHeader(fi, "")
	if err != nil {
		return nil, trace.Wrap(err)
	}
	item := &Item{
		Header: *fiHeader,
	}
	item.Name = localPath
	if !fi.IsDir() {
		item.Data, err = os.Open(path)
		if err != nil {
			return nil, trace.Wrap(err)
		}
	}
	return item, nil
}

type Item struct {
	tar.Header

	Data io.ReadCloser
}

func CreateMemArchive(items []*Item) (*bytes.Buffer, error) {
	buf := &bytes.Buffer{}
	archive := NewTarAppender(buf)

	if err := archive.Add(items...); err != nil {
		return nil, trace.Wrap(err)
	}
	if err := archive.Close(); err != nil {
		return nil, trace.Wrap(err)
	}

	return buf, nil
}

func MustCreateMemArchive(items []*Item) *bytes.Buffer {
	r, err := CreateMemArchive(items)
	if err != nil {
		panic(err)
	}
	return r
}

var Abort = errors.New("abort iteration")

func extractFile(tarball *tar.Reader, header *tar.Header, dir string) error {
	switch header.Typeflag {
	case tar.TypeDir:
		return withDir(filepath.Join(dir, header.Name), nil)
	case tar.TypeBlock, tar.TypeChar, tar.TypeReg, tar.TypeRegA, tar.TypeFifo:
		return writeFile(filepath.Join(dir, header.Name), tarball, header.FileInfo().Mode())
	case tar.TypeLink:
		return writeHardLink(filepath.Join(dir, header.Name), filepath.Join(dir, header.Linkname))
	case tar.TypeSymlink:
		return writeSymbolicLink(filepath.Join(dir, header.Name), header.Linkname)
	default:
		log.Warnf("unsupported type flag %v for %v", header.Typeflag, header.Name)
	}
	return nil
}

func writeFile(path string, r io.Reader, mode os.FileMode) error {
	err := withDir(path, func() error {
		out, err := os.Create(path)
		if err != nil {
			return trace.ConvertSystemError(err)
		}
		defer out.Close()

		err = out.Chmod(mode)
		if err != nil {
			return trace.ConvertSystemError(err)
		}

		_, err = io.Copy(out, r)
		return trace.Wrap(err)
	})
	return trace.Wrap(err)
}

func writeSymbolicLink(path string, target string) error {
	err := withDir(path, func() error {
		err := os.Symlink(target, path)
		return trace.ConvertSystemError(err)
	})
	return trace.Wrap(err)
}

func writeHardLink(path string, target string) error {
	err := withDir(path, func() error {
		err := os.Link(target, path)
		return trace.ConvertSystemError(err)
	})
	return trace.Wrap(err)
}

func withDir(path string, fn func() error) error {
	err := os.MkdirAll(filepath.Dir(path), defaults.SharedDirMask)
	if err != nil {
		return trace.ConvertSystemError(err)
	}

	if fn == nil {
		return nil
	}
	err = fn()
	return trace.Wrap(err)
}
/*
Copyright 2018 Gravitational, Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package archive

import (
	"archive/tar"
	"bytes"
	"errors"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/gravitational/gravity/lib/defaults"

	dockerarchive "github.com/docker/docker/pkg/archive"
	"github.com/gravitational/trace"
	log "github.com/sirupsen/logrus"
)

// DefaultOptions returns an options object with sensible defaults
func DefaultOptions() *dockerarchive.TarOptions {
	options := &dockerarchive.TarOptions{NoLchown: true}
	return options
}

// CompressDirectory compresses the directory given with dir, using writer as a sink
// for the archive
func CompressDirectory(dir string, writer io.Writer, items ...*Item) error {
	archive := NewTarAppender(writer)
	defer archive.Close()

	if err := archive.Add(items...); err != nil {
		return trace.Wrap(err, "failed to write tarball: %v", err.Error())
	}
	if err := filepath.Walk(dir, func(path string, fi os.FileInfo, err error) error {
		if err != nil {
			return trace.Wrap(err)
		}
		localPath, err := filepath.Rel(dir, path)
		if err != nil {
			return trace.Wrap(err)
		}
		if localPath == "." {
			// Skip current directory item
			return nil
		}
		item, err := ItemFromFile(localPath, path, fi)
		if err != nil {
			return trace.Wrap(err)
		}
		return trace.Wrap(archive.Add(item))
	}); err != nil {
		return trace.Wrap(err, "failed to compress directory %q", dir)
	}
	return nil
}

// Unpack unpacks the specified tarball to a temporary directory and returns
// the directory where it was unpacked
func Unpack(path string) (unpackedDir string, err error) {
	file, err := os.Open(path)
	if err != nil {
		return "", trace.ConvertSystemError(err)
	}
	defer file.Close()
	tmp, err := ioutil.TempDir("", "")
	if err != nil {
		return "", trace.ConvertSystemError(err)
	}
	err = Extract(file, tmp)
	if err != nil {
		return "", trace.Wrap(err)
	}
	return tmp, nil
}

// Extract extracts the contents of the specified tarball under dir.
// The resulting files and directories are created using the current user context.
func Extract(r io.Reader, dir string) error {
	tarball := tar.NewReader(r)

	for {
		header, err := tarball.Next()
		if err == io.EOF {
			break
		} else if err != nil {
			return trace.Wrap(err)
		}

		if err := extractFile(tarball, header, dir); err != nil {
			return trace.Wrap(err)
		}
	}
	return nil
}

// HasFile returns nil if the specified tarball contains specified file
func HasFile(tarballPath, filename string) error {
	file, err := os.Open(tarballPath)
	if err != nil {
		return trace.ConvertSystemError(err)
	}
	defer file.Close()
	var hasFile bool
	err = TarGlob(tar.NewReader(file), ".", []string{filename},
		func(match string, file io.Reader) error {
			hasFile = true
			return Abort
		})
	if err != nil {
		if trace.Unwrap(err) == tar.ErrHeader {
			return trace.BadParameter("file %v does not appear to be a valid tarball",
				tarballPath)
		}
		return trace.Wrap(err)
	}
	if !hasFile {
		return trace.NotFound("tarball %v does not contain file %v",
			tarballPath, filename)
	}
	return nil
}

// TarGlob iterates the contents of the specified tarball and invokes handler
// for each file matching the list of specified patterns.
// If the handler returns a special Abort error, iteration will be aborted without errors.
func TarGlob(source *tar.Reader, dir string, patterns []string, handler func(match string, file io.Reader) error) (err error) {
	for {
		var hdr *tar.Header
		hdr, err = source.Next()
		if err == io.EOF {
			err = nil
			break
		}
		if err != nil {
			return trace.Wrap(err)
		}
		if hdr.FileInfo().IsDir() {
			continue
		}
		for _, pattern := range patterns {
			relpath, err := filepath.Rel(dir, hdr.Name)
			if err == nil {
				matched, _ := filepath.Match(pattern, filepath.Base(relpath))
				if matched {
					if err = handler(relpath, source); err != nil {
						if trace.Unwrap(err) == Abort {
							return nil
						}
						return trace.Wrap(err)
					}
				}
			}
		}
	}
	return nil
}

// TarAppender wraps a tar writer and can append items to it
type TarAppender struct {
	tw *tar.Writer
}

// NewTarAppender creates a new tar appender writing to w
func NewTarAppender(w io.Writer) *TarAppender {
	return &TarAppender{tar.NewWriter(w)}
}

// Add adds the specified items to the underlined archive
func (r *TarAppender) Add(items ...*Item) (err error) {
	defer func() {
		for _, item := range items {
			if item.Data != nil {
				item.Data.Close()
			}
		}
	}()
	for _, item := range items {
		if item.ModTime.IsZero() {
			item.ModTime = time.Now()
		}
		if err = r.tw.WriteHeader(&item.Header); err != nil {
			return trace.Wrap(err)
		}
		if item.Data == nil {
			continue
		}
		if _, err = io.Copy(r.tw, item.Data); err != nil {
			return trace.Wrap(err)
		}
	}
	return nil
}

// Close closes the underlined archive
func (r *TarAppender) Close() error {
	return r.tw.Close()
}

// ItemFromString creates an Item from given string
func ItemFromString(path, value string) *Item {
	return ItemFromStringMode(path, value, defaults.SharedExecutableMask)
}

// DirItem creates a new virtual directory item
func DirItem(path string) *Item {
	return &Item{
		Header: tar.Header{
			Name:     path + "/",
			Typeflag: tar.TypeDir,
			Mode:     defaults.SharedDirMask,
			Uid:      defaults.ArchiveUID,
			Gid:      defaults.ArchiveGID,
		},
	}
}

// ItemFromStringMode creates a new Item from a given string with the provided permissions
func ItemFromStringMode(path, value string, mode int64) *Item {
	return ItemFromStream(path, ioutil.NopCloser(strings.NewReader(value)), int64(len(value)), mode)
}

// ItemFromStream creates an Item from given io.ReadCloser
func ItemFromStream(path string, rc io.ReadCloser, size, mode int64) *Item {
	return &Item{
		Header: tar.Header{
			Name: path,
			Size: size,
			Mode: mode,
			Uid:  defaults.ArchiveUID,
			Gid:  defaults.ArchiveGID,
		},
		Data: rc,
	}
}

// ItemFromFile creates an Item from the specified file
func ItemFromFile(localPath, path string, fi os.FileInfo) (*Item, error) {
	fiHeader, err := tar.FileInfoHeader(fi, "")
	if err != nil {
		return nil, trace.Wrap(err)
	}
	item := &Item{
		Header: *fiHeader,
	}
	item.Name = localPath
	if !fi.IsDir() {
		item.Data, err = os.Open(path)
		if err != nil {
			return nil, trace.Wrap(err)
		}
	}
	return item, nil
}

// Item defines a unit of compression
type Item struct {
	tar.Header
	// Data is the data payload
	Data io.ReadCloser
}

// CreateMemArchive creates in-memory archive from archive items
func CreateMemArchive(items []*Item) (*bytes.Buffer, error) {
	buf := &bytes.Buffer{}
	archive := NewTarAppender(buf)

	if err := archive.Add(items...); err != nil {
		return nil, trace.Wrap(err)
	}
	if err := archive.Close(); err != nil {
		return nil, trace.Wrap(err)
	}

	return buf, nil
}

// MustCreateMemArchive creates in-memory archive from archive items
func MustCreateMemArchive(items []*Item) *bytes.Buffer {
	r, err := CreateMemArchive(items)
	if err != nil {
		panic(err)
	}
	return r
}

// Abort is a special error value used to abort an iteration
var Abort = errors.New("abort iteration")

// extractFile extracts a single file or directory from tarball into dir.
// Uses header to determine the type of item to create
// Based on https://github.com/mholt/archiver
func extractFile(tarball *tar.Reader, header *tar.Header, dir string) error {
	switch header.Typeflag {
	case tar.TypeDir:
		return withDir(filepath.Join(dir, header.Name), nil)
	case tar.TypeBlock, tar.TypeChar, tar.TypeReg, tar.TypeRegA, tar.TypeFifo:
		return writeFile(filepath.Join(dir, header.Name), tarball, header.FileInfo().Mode())
	case tar.TypeLink:
		return writeHardLink(filepath.Join(dir, header.Name), filepath.Join(dir, header.Linkname))
	case tar.TypeSymlink:
		return writeSymbolicLink(filepath.Join(dir, header.Name), header.Linkname)
	default:
		log.Warnf("unsupported type flag %v for %v", header.Typeflag, header.Name)
	}
	return nil
}

func writeFile(path string, r io.Reader, mode os.FileMode) error {
	err := withDir(path, func() error {
		out, err := os.Create(path)
		if err != nil {
			return trace.ConvertSystemError(err)
		}
		defer out.Close()

		err = out.Chmod(mode)
		if err != nil {
			return trace.ConvertSystemError(err)
		}

		_, err = io.Copy(out, r)
		return trace.Wrap(err)
	})
	return trace.Wrap(err)
}

func writeSymbolicLink(path string, target string) error {
	err := withDir(path, func() error {
		err := os.Symlink(target, path)
		return trace.ConvertSystemError(err)
	})
	return trace.Wrap(err)
}

func writeHardLink(path string, target string) error {
	err := withDir(path, func() error {
		err := os.Link(target, path)
		return trace.ConvertSystemError(err)
	})
	return trace.Wrap(err)
}

func withDir(path string, fn func() error) error {
	err := os.MkdirAll(filepath.Dir(path), defaults.SharedDirMask)
	if err != nil {
		return trace.ConvertSystemError(err)
	}

	if fn == nil {
		return nil
	}
	err = fn()
	return trace.Wrap(err)
}
