// Package archive streams a directory through tar+zstd (and back).
//
// Compression level is fixed at zstd level 3 — same as the existing shell
// (`zstd -T0 -3`). The encoder is wrapped around an io.Writer (typically an
// S3 multipart Uploader's PipeWriter), so the whole pipeline stays
// streaming end-to-end with no on-disk tempfile.
package archive

import (
	"archive/tar"
	"errors"
	"fmt"
	"io"
	"io/fs"
	"os"
	"path/filepath"
	"sort"
	"strings"

	"github.com/klauspost/compress/zstd"
)

// EncodeDir walks root in sorted relative-path order and writes a
// tar.zst stream into w. Symlinks are preserved (stored as tar typelink),
// devices/sockets/fifos are skipped.
func EncodeDir(w io.Writer, root string) error {
	zw, err := zstd.NewWriter(w, zstd.WithEncoderLevel(zstd.SpeedDefault))
	if err != nil {
		return fmt.Errorf("zstd writer: %w", err)
	}
	defer zw.Close()
	tw := tar.NewWriter(zw)

	var paths []string
	err = filepath.WalkDir(root, func(p string, d fs.DirEntry, err error) error {
		if err != nil {
			return err
		}
		if p == root {
			return nil
		}
		paths = append(paths, p)
		return nil
	})
	if err != nil {
		return fmt.Errorf("walk %s: %w", root, err)
	}
	sort.Strings(paths)

	for _, p := range paths {
		if err := writeOne(tw, root, p); err != nil {
			return err
		}
	}
	if err := tw.Close(); err != nil {
		return fmt.Errorf("close tar: %w", err)
	}
	if err := zw.Close(); err != nil {
		return fmt.Errorf("close zstd: %w", err)
	}
	return nil
}

func writeOne(tw *tar.Writer, root, p string) error {
	rel, err := filepath.Rel(root, p)
	if err != nil {
		return err
	}
	rel = filepath.ToSlash(rel)

	lst, err := os.Lstat(p)
	if err != nil {
		return fmt.Errorf("lstat %s: %w", p, err)
	}

	var link string
	if lst.Mode()&os.ModeSymlink != 0 {
		link, err = os.Readlink(p)
		if err != nil {
			return fmt.Errorf("readlink %s: %w", p, err)
		}
	}
	hdr, err := tar.FileInfoHeader(lst, link)
	if err != nil {
		// Skip unsupported file types (devices, sockets, fifos).
		return nil
	}
	hdr.Name = rel
	if lst.IsDir() && !strings.HasSuffix(hdr.Name, "/") {
		hdr.Name += "/"
	}
	if err := tw.WriteHeader(hdr); err != nil {
		return fmt.Errorf("tar header %s: %w", rel, err)
	}
	if !lst.Mode().IsRegular() {
		return nil
	}
	f, err := os.Open(p)
	if err != nil {
		return fmt.Errorf("open %s: %w", p, err)
	}
	defer f.Close()
	if _, err := io.Copy(tw, f); err != nil {
		return fmt.Errorf("copy %s: %w", rel, err)
	}
	return nil
}

// DecodeDir extracts a tar.zst stream from r into dest. dest must exist.
// Any entry whose normalized path escapes dest is rejected.
//
// Read-only entries are the normal case, not an exotic one: the Go module
// cache is 0555 directories and 0444 files throughout. So directories are
// created writable and their recorded mode is applied at the very end,
// and files are written writable and chmod'ed after the copy. Restoring
// the archived mode as we go would lock the extraction out of the tree it
// is still filling in.
func DecodeDir(r io.Reader, dest string) error {
	absDest, err := filepath.Abs(dest)
	if err != nil {
		return fmt.Errorf("abs %s: %w", dest, err)
	}
	zr, err := zstd.NewReader(r)
	if err != nil {
		return fmt.Errorf("zstd reader: %w", err)
	}
	defer zr.Close()
	tr := tar.NewReader(zr)

	// Directory modes, in archive order, applied once everything is in
	// place. Deepest-first on the way back out.
	var dirs []dirMode

	for {
		hdr, err := tr.Next()
		if errors.Is(err, io.EOF) {
			return restoreDirModes(dirs)
		}
		if err != nil {
			return fmt.Errorf("tar next: %w", err)
		}
		target, err := safeJoin(absDest, hdr.Name)
		if err != nil {
			return err
		}
		mode := fs.FileMode(hdr.Mode) & 0o7777
		switch hdr.Typeflag {
		case tar.TypeDir:
			if err := os.MkdirAll(target, mode|0o700); err != nil {
				return fmt.Errorf("mkdir %s: %w", target, err)
			}
			dirs = append(dirs, dirMode{path: target, mode: mode})
		case tar.TypeReg, tar.TypeRegA:
			if err := os.MkdirAll(filepath.Dir(target), 0o755); err != nil {
				return fmt.Errorf("mkdir parent of %s: %w", target, err)
			}
			f, err := createFile(target, mode)
			if err != nil {
				return fmt.Errorf("create %s: %w", target, err)
			}
			if _, err := io.Copy(f, tr); err != nil {
				f.Close()
				return fmt.Errorf("write %s: %w", target, err)
			}
			if err := f.Close(); err != nil {
				return err
			}
			if err := os.Chmod(target, mode); err != nil {
				return fmt.Errorf("chmod %s: %w", target, err)
			}
		case tar.TypeSymlink:
			_ = os.Remove(target)
			if err := os.Symlink(hdr.Linkname, target); err != nil {
				return fmt.Errorf("symlink %s -> %s: %w", target, hdr.Linkname, err)
			}
		default:
			// Skip unknown entry types silently.
		}
	}
}

// dirMode is a directory whose archived mode is applied after extraction.
type dirMode struct {
	path string
	mode fs.FileMode
}

// restoreDirModes applies the archived modes deepest-first: entries arrive
// parent-before-child, so walking back tightens a directory only after
// everything under it is already done.
func restoreDirModes(dirs []dirMode) error {
	for i := len(dirs) - 1; i >= 0; i-- {
		if err := os.Chmod(dirs[i].path, dirs[i].mode); err != nil {
			return fmt.Errorf("chmod %s: %w", dirs[i].path, err)
		}
	}
	return nil
}

// createFile opens target for writing, forcing the owner-write bit on so a
// read-only archived mode doesn't stop us writing the content. The caller
// chmods to the archived mode afterwards. A read-only file left behind by
// an earlier extraction is replaced rather than treated as an error —
// re-restoring over a populated module cache is a normal CI situation.
func createFile(target string, mode fs.FileMode) (*os.File, error) {
	f, err := os.OpenFile(target, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, mode|0o200)
	if err == nil || !errors.Is(err, fs.ErrPermission) {
		return f, err
	}
	if rmErr := os.Remove(target); rmErr != nil {
		// The parent directory may itself be read-only; widen it and
		// retry once before giving up with the original error.
		parent := filepath.Dir(target)
		st, stErr := os.Stat(parent)
		if stErr != nil {
			return nil, err
		}
		if chErr := os.Chmod(parent, st.Mode().Perm()|0o700); chErr != nil {
			return nil, err
		}
		if rmErr := os.Remove(target); rmErr != nil {
			return nil, err
		}
	}
	return os.OpenFile(target, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, mode|0o200)
}

// safeJoin returns base/sub, rejecting any sub that is absolute or contains
// ".." segments that would resolve outside base. We reject rather than
// silently re-root so malicious tarballs surface as an error.
func safeJoin(base, sub string) (string, error) {
	if filepath.IsAbs(sub) || strings.HasPrefix(sub, "/") {
		return "", fmt.Errorf("absolute tar entry: %q", sub)
	}
	clean := filepath.Clean(sub)
	if clean == ".." || strings.HasPrefix(clean, ".."+string(filepath.Separator)) {
		return "", fmt.Errorf("tar entry escapes destination: %q", sub)
	}
	return filepath.Join(base, clean), nil
}