package archive

import (
	"bytes"
	"fmt"
	"io"
	"os/exec"
	"strings"

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

// BuildSourceArchive captures the working tree as `tar.zst` bytes.
//
// When the working tree is clean it archives HEAD. When it's dirty
// (modified or staged tracked files), it archives a temporary stash
// object so the dirty state is what ships — callers that want
// HEAD-only behaviour should reject dirty trees before calling.
// Untracked files are never included regardless of state.
func BuildSourceArchive(repoDir string) ([]byte, error) {
	stashCmd := exec.Command("git", "stash", "create")
	stashCmd.Dir = repoDir
	var stashOut, stashErr bytes.Buffer
	stashCmd.Stdout = &stashOut
	stashCmd.Stderr = &stashErr
	if err := stashCmd.Run(); err != nil {
		return nil, fmt.Errorf("git stash create: %v: %s", err, stashErr.String())
	}
	treeish := "HEAD"
	if sha := strings.TrimSpace(stashOut.String()); sha != "" {
		treeish = sha
	}

	cmd := exec.Command("git", "archive", "--format=tar", treeish)
	cmd.Dir = repoDir
	var tarOut, stderr bytes.Buffer
	cmd.Stdout = &tarOut
	cmd.Stderr = &stderr
	if err := cmd.Run(); err != nil {
		return nil, fmt.Errorf("git archive: %v: %s", err, stderr.String())
	}

	var compressed bytes.Buffer
	enc, err := zstd.NewWriter(&compressed, zstd.WithEncoderLevel(zstd.SpeedDefault))
	if err != nil {
		return nil, err
	}
	if _, err := io.Copy(enc, &tarOut); err != nil {
		_ = enc.Close()
		return nil, err
	}
	if err := enc.Close(); err != nil {
		return nil, err
	}
	return compressed.Bytes(), nil
}