package main

import (
	"archive/tar"
	"archive/zip"
	"bytes"
	"compress/bzip2"
	"compress/gzip"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

// Single-file compression that Go does not carry in its standard library. Each
// entry lists the tools that can expand it, in order of preference; the first
// one found wins. All of them write to stdout with the flags given.
var external = map[string]struct {
	tools []string
	args  []string
}{
	"xz":   {[]string{"xz", "unxz"}, []string{"-dc"}},
	"lzma": {[]string{"xz", "unxz", "lzma"}, []string{"-dc"}},
	"zst":  {[]string{"zstd", "unzstd"}, []string{"-dcq"}},
	"lz4":  {[]string{"lz4", "unlz4"}, []string{"-dcq"}},
}

type decomp struct {
	prog string
	args []string
}

// decompressor reports how to expand ext, or nil if this machine cannot.
// gz and bz2 are handled in-process and need no helper at all.
func decompressor(ext string) *decomp {
	switch ext {
	case "gz", "bz2":
		return &decomp{} // stdlib, no external program
	}
	e, ok := external[ext]
	if !ok {
		return nil
	}
	for _, t := range e.tools {
		if p, err := exec.LookPath(t); err == nil {
			return &decomp{prog: p, args: e.args}
		}
	}
	return nil
}

func sevenZip() string {
	for _, t := range []string{"7zz", "7z", "7za"} {
		if p, err := exec.LookPath(t); err == nil {
			return p
		}
	}
	return ""
}

// reader wraps f in whatever ext needs, shelling out only when the standard
// library has no codec for it.
func decompressReader(ext, path string, f *os.File) (io.Reader, func() error, error) {
	nop := func() error { return nil }
	switch ext {
	case "gz":
		zr, err := gzip.NewReader(f)
		if err != nil {
			return nil, nil, err
		}
		return zr, zr.Close, nil
	case "bz2":
		return bzip2.NewReader(f), nop, nil
	}

	d := decompressor(ext)
	if d == nil {
		return nil, nil, fmt.Errorf("cannot unpack .%s - install %s or pick another asset with --asset/--pattern",
			ext, strings.Join(external[ext].tools, "/"))
	}
	cmd := exec.Command(d.prog, append(append([]string{}, d.args...), path)...)
	cmd.Stdin = nil
	cmd.Stderr = os.Stderr
	out, err := cmd.StdoutPipe()
	if err != nil {
		return nil, nil, err
	}
	if err := cmd.Start(); err != nil {
		return nil, nil, err
	}
	return out, func() error {
		io.Copy(io.Discard, out) // drain, or the child gets EPIPE
		if err := cmd.Wait(); err != nil {
			return fmt.Errorf("failed to decompress %s (%s: %w)", filepath.Base(path), filepath.Base(d.prog), err)
		}
		return nil
	}, nil
}

// extractIfArchive unpacks into a fresh directory below tmpdir and returns its
// path, or "" if the asset is a bare binary (then the download is the binary).
func extractIfArchive(file, tmpdir string) (string, error) {
	lc := strings.ToLower(filepath.Base(file))
	out := filepath.Join(tmpdir, "x")

	switch {
	case tarballRe.MatchString(lc):
		if err := os.Mkdir(out, 0o755); err != nil {
			return "", err
		}
		return out, untar(file, out, tarCodec(lc))

	case strings.HasSuffix(lc, ".zip"):
		if err := os.Mkdir(out, 0o755); err != nil {
			return "", err
		}
		return out, unzip(file, out)

	case strings.HasSuffix(lc, ".7z"):
		z := sevenZip()
		if z == "" {
			return "", fmt.Errorf("7z not found, but needed for %s", lc)
		}
		if err := os.Mkdir(out, 0o755); err != nil {
			return "", err
		}
		cmd := exec.Command(z, "x", "-y", "-o"+out, file)
		cmd.Stdout, cmd.Stderr = io.Discard, os.Stderr
		if err := cmd.Run(); err != nil {
			return "", fmt.Errorf("failed to extract %s: %w", lc, err)
		}
		return out, nil
	}

	// A single compressed file, e.g. restic_0.19.1_darwin_arm64.bz2 - the
	// binary itself, just squeezed.
	if m := singleRe.FindStringSubmatch(lc); m != nil {
		ext := m[1]
		if err := os.Mkdir(out, 0o755); err != nil {
			return "", err
		}
		stem := strings.TrimSuffix(filepath.Base(file), filepath.Ext(file))
		plain := filepath.Join(out, stem)
		if err := decompressFile(ext, file, plain); err != nil {
			return "", err
		}
		return out, nil
	}
	return "", nil // bare binary
}

// tarCodec maps a tarball suffix to the compression wrapped around the tar.
func tarCodec(lc string) string {
	switch {
	case strings.HasSuffix(lc, ".tar.gz"), strings.HasSuffix(lc, ".tgz"):
		return "gz"
	case strings.HasSuffix(lc, ".tar.bz2"), strings.HasSuffix(lc, ".tbz"):
		return "bz2"
	case strings.HasSuffix(lc, ".tar.xz"), strings.HasSuffix(lc, ".txz"):
		return "xz"
	case strings.HasSuffix(lc, ".tar.zst"), strings.HasSuffix(lc, ".tzst"):
		return "zst"
	}
	return "" // plain .tar
}

func decompressFile(ext, src, dst string) error {
	f, err := os.Open(src)
	if err != nil {
		return err
	}
	defer f.Close()

	r, closer, err := decompressReader(ext, src, f)
	if err != nil {
		return err
	}
	out, err := os.Create(dst)
	if err != nil {
		return err
	}
	defer out.Close()
	if _, err := io.Copy(out, r); err != nil {
		return err
	}
	return closer()
}

func untar(file, dest, codec string) error {
	f, err := os.Open(file)
	if err != nil {
		return err
	}
	defer f.Close()

	var r io.Reader = f
	closer := func() error { return nil }
	if codec != "" {
		r, closer, err = decompressReader(codec, file, f)
		if err != nil {
			return err
		}
	}

	tr := tar.NewReader(r)
	for {
		h, err := tr.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			return fmt.Errorf("failed to extract %s: %w", filepath.Base(file), err)
		}
		path, ok := safeJoin(dest, h.Name)
		if !ok {
			verbose("skipping %s (path escapes the archive)", h.Name)
			continue
		}
		switch h.Typeflag {
		case tar.TypeDir:
			if err := os.MkdirAll(path, 0o755); err != nil {
				return err
			}
		case tar.TypeReg:
			if err := writeFileFrom(path, tr, os.FileMode(h.Mode).Perm()); err != nil {
				return err
			}
		default:
			// Symlinks and devices are never the binary we are after.
			verbose("skipping %s (not a regular file)", h.Name)
		}
	}
	return closer()
}

func unzip(file, dest string) error {
	zr, err := zip.OpenReader(file)
	if err != nil {
		return fmt.Errorf("failed to extract %s: %w", filepath.Base(file), err)
	}
	defer zr.Close()

	for _, f := range zr.File {
		path, ok := safeJoin(dest, f.Name)
		if !ok {
			verbose("skipping %s (path escapes the archive)", f.Name)
			continue
		}
		if f.FileInfo().IsDir() {
			if err := os.MkdirAll(path, 0o755); err != nil {
				return err
			}
			continue
		}
		if !f.Mode().IsRegular() {
			continue
		}
		rc, err := f.Open()
		if err != nil {
			return err
		}
		err = writeFileFrom(path, rc, f.Mode().Perm())
		rc.Close()
		if err != nil {
			return err
		}
	}
	return nil
}

// safeJoin keeps an archive from writing outside its own directory.
func safeJoin(dest, name string) (string, bool) {
	clean := filepath.Clean(filepath.FromSlash(name))
	if filepath.IsAbs(clean) || clean == ".." || strings.HasPrefix(clean, ".."+string(os.PathSeparator)) {
		return "", false
	}
	return filepath.Join(dest, clean), true
}

func writeFileFrom(path string, r io.Reader, mode os.FileMode) error {
	if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
		return err
	}
	if mode == 0 {
		mode = 0o644
	}
	f, err := os.OpenFile(path, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, mode)
	if err != nil {
		return err
	}
	defer f.Close()
	_, err = io.Copy(f, r)
	return err
}

// findInTree looks for a file called name; failing that, for the single
// executable in the tree.
func findInTree(root, name string) string {
	var files []string
	filepath.WalkDir(root, func(p string, d os.DirEntry, err error) error {
		if err == nil && d.Type().IsRegular() {
			files = append(files, p)
		}
		return nil
	})

	for _, p := range files {
		if b := filepath.Base(p); b == name || b == name+".exe" {
			return p
		}
	}
	var execs []string
	for _, p := range files {
		if looksExecutable(p) {
			execs = append(execs, p)
		}
	}
	if len(execs) == 1 {
		return execs[0]
	}
	verbose("archive contains %d files, %d of them executable", len(files), len(execs))
	return ""
}

// Magic bytes beat guessing by file size: ELF, Mach-O (incl. fat), PE, script.
func looksExecutable(path string) bool {
	f, err := os.Open(path)
	if err != nil {
		return false
	}
	defer f.Close()

	m := make([]byte, 4)
	n, _ := io.ReadFull(f, m)
	if n < 2 {
		return false
	}
	m = m[:n]
	switch {
	case bytes.HasPrefix(m, []byte("\x7fELF")):
		return true
	case bytes.HasPrefix(m, []byte{0xcf, 0xfa, 0xed, 0xfe}), // Mach-O 64
		bytes.HasPrefix(m, []byte{0xce, 0xfa, 0xed, 0xfe}), // Mach-O 32
		bytes.HasPrefix(m, []byte{0xca, 0xfe, 0xba, 0xbe}), // fat
		bytes.HasPrefix(m, []byte{0xbe, 0xba, 0xfe, 0xca}):
		return true
	case bytes.HasPrefix(m, []byte("MZ")):
		return true
	case bytes.HasPrefix(m, []byte("#!")):
		return true
	}
	return false
}