ai-workflow-skill/internal/cli/orch/worktree.go

package orch

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"ai-workflow-skill/internal/protocol"
	"ai-workflow-skill/internal/store"

	"github.com/spf13/cobra"
)

func newDispatchWorkspacePreparer(cmd *cobra.Command, opts dispatchOptions) store.DispatchWorkspacePreparer {
	ctx := cmd.Context()

	return func(task store.Task, attemptNo int) (store.DispatchWorkspace, func(), error) {
		effectiveOpts, useWorktree := resolveDispatchWorktreeOptions(task, opts)
		if !useWorktree {
			return store.DispatchWorkspace{}, func() {}, nil
		}
		return provisionDispatchWorkspace(ctx, effectiveOpts, task, attemptNo)
	}
}

func newAttemptReuseWorkspacePreparer(cmd *cobra.Command, task store.Task, attempt *store.TaskAttempt) store.DispatchWorkspacePreparer {
	if attempt == nil || attempt.WorktreePath == "" {
		return nil
	}

	workspaceRoot, ok := deriveWorkspaceRootFromAttempt(task.RunID, task.TaskID, attempt.WorktreePath)
	if !ok {
		return nil
	}

	baseRef := attempt.BaseRef
	if strings.TrimSpace(baseRef) == "" {
		baseRef = attempt.BaseCommit
	}

	opts := dispatchOptions{
		repoPath:       attempt.WorktreePath,
		workspaceRoot:  workspaceRoot,
		strictWorktree: true,
		baseRef:        baseRef,
	}

	return newDispatchWorkspacePreparer(cmd, opts)
}

func dispatchUsesWorktree(opts dispatchOptions) bool {
	return strings.TrimSpace(opts.workspaceRoot) != "" ||
		opts.strictWorktree
}

func resolveDispatchWorktreeOptions(task store.Task, opts dispatchOptions) (dispatchOptions, bool) {
	if dispatchUsesWorktree(opts) {
		return opts, true
	}
	if !taskLooksLikeCodeWork(task) {
		return opts, false
	}

	auto := opts
	auto.strictWorktree = true
	return auto, true
}

func taskLooksLikeCodeWork(task store.Task) bool {
	if acceptanceJSONLooksCodeLike(task.AcceptanceJSON) {
		return true
	}
	return roleLooksCodeLike(task.DefaultTo)
}

func acceptanceJSONLooksCodeLike(raw json.RawMessage) bool {
	if len(raw) == 0 {
		return false
	}

	var value any
	if err := json.Unmarshal(raw, &value); err != nil {
		return false
	}
	return acceptanceValueLooksCodeLike(value)
}

func acceptanceValueLooksCodeLike(value any) bool {
	switch typed := value.(type) {
	case map[string]any:
		for key, raw := range typed {
			lowerKey := strings.ToLower(strings.TrimSpace(key))
			switch lowerKey {
			case "code", "code_task", "writes_code", "worktree":
				if boolValue, ok := raw.(bool); ok && boolValue {
					return true
				}
			case "kind", "task_type", "mode", "type":
				if stringValue, ok := raw.(string); ok && isCodeLikeMarker(stringValue) {
					return true
				}
			}
			if acceptanceValueLooksCodeLike(raw) {
				return true
			}
		}
	case []any:
		for _, item := range typed {
			if acceptanceValueLooksCodeLike(item) {
				return true
			}
		}
	case string:
		return isCodeLikeMarker(typed)
	}
	return false
}

func roleLooksCodeLike(role string) bool {
	role = strings.ToLower(strings.TrimSpace(role))
	if role == "" {
		return false
	}

	for _, token := range splitIdentifierTokens(role) {
		switch token {
		case "backend", "frontend", "front", "admin", "ui", "fullstack", "foundation", "db", "database", "mobile", "ios", "android", "web", "platform", "infra", "api":
			return true
		}
	}
	return false
}

func isCodeLikeMarker(value string) bool {
	value = strings.ToLower(strings.TrimSpace(value))
	switch value {
	case "code", "code_task", "code-task", "code-change", "code_change", "implementation", "patch", "diff", "repo":
		return true
	default:
		return false
	}
}

func splitIdentifierTokens(value string) []string {
	return strings.FieldsFunc(value, func(r rune) bool {
		return !((r >= 'a' && r <= 'z') || (r >= '0' && r <= '9'))
	})
}

func provisionDispatchWorkspace(ctx context.Context, opts dispatchOptions, task store.Task, attemptNo int) (store.DispatchWorkspace, func(), error) {
	repoRoot, err := resolveRepoRoot(ctx, opts.repoPath)
	if err != nil {
		return store.DispatchWorkspace{}, nil, err
	}

	workspaceRoot := resolveWorkspaceRoot(repoRoot, opts.workspaceRoot)
	if err := ensureWorkspaceRootIgnored(repoRoot, workspaceRoot); err != nil {
		return store.DispatchWorkspace{}, nil, err
	}

	baseRef, baseCommit, err := resolveDispatchBase(ctx, repoRoot, workspaceRoot, opts.baseRef, opts.strictWorktree)
	if err != nil {
		return store.DispatchWorkspace{}, nil, err
	}

	branchName := buildAttemptBranchName(task.RunID, task.TaskID, attemptNo)
	worktreePath := buildAttemptWorktreePath(workspaceRoot, task.RunID, task.TaskID, attemptNo)

	if err := os.MkdirAll(filepath.Dir(worktreePath), 0o755); err != nil {
		return store.DispatchWorkspace{}, nil, fmt.Errorf("create worktree parent dir: %w", err)
	}
	if _, err := os.Stat(worktreePath); err == nil {
		return store.DispatchWorkspace{}, nil, fmt.Errorf("%w: worktree path already exists: %s", store.ErrInvalidState, worktreePath)
	} else if err != nil && !os.IsNotExist(err) {
		return store.DispatchWorkspace{}, nil, fmt.Errorf("stat worktree path: %w", err)
	}

	if _, _, err := runGit(ctx, repoRoot, "worktree", "add", "-b", branchName, worktreePath, baseCommit); err != nil {
		return store.DispatchWorkspace{}, nil, err
	}

	cleanup := func() {
		_, _, _ = runGit(context.Background(), repoRoot, "worktree", "remove", "--force", worktreePath)
		_, _, _ = runGit(context.Background(), repoRoot, "branch", "-D", branchName)
		_ = os.RemoveAll(worktreePath)
	}

	return store.DispatchWorkspace{
		BaseRef:         baseRef,
		BaseCommit:      baseCommit,
		BranchName:      branchName,
		WorktreePath:    worktreePath,
		WorkspaceStatus: "created",
	}, cleanup, nil
}

func resolveRepoRoot(ctx context.Context, repoPath string) (string, error) {
	startPath := strings.TrimSpace(repoPath)
	if startPath == "" {
		var err error
		startPath, err = os.Getwd()
		if err != nil {
			return "", fmt.Errorf("get current working directory: %w", err)
		}
	}

	absPath, err := filepath.Abs(startPath)
	if err != nil {
		return "", fmt.Errorf("resolve repo path: %w", err)
	}

	if _, _, err := runGit(ctx, absPath, "rev-parse", "--show-toplevel"); err != nil {
		return "", protocol.InvalidInput("repo-path must point to a Git worktree", err)
	}

	commonDir, err := resolveCommonGitDir(ctx, absPath)
	if err != nil {
		return "", protocol.InvalidInput("repo-path must point to a Git worktree", err)
	}
	return filepath.Dir(commonDir), nil
}

func resolveDispatchBase(ctx context.Context, repoRoot, workspaceRoot, requestedBaseRef string, strict bool) (string, string, error) {
	baseRef := strings.TrimSpace(requestedBaseRef)
	if baseRef != "" {
		baseCommit, err := resolveCommit(ctx, repoRoot, baseRef)
		if err != nil {
			return "", "", protocol.InvalidInput("base-ref must resolve to a commit", err)
		}
		return baseRef, baseCommit, nil
	}

	if strict {
		dirty, err := repoHasUncommittedChanges(ctx, repoRoot, workspaceRoot)
		if err != nil {
			return "", "", err
		}
		if dirty {
			return "", "", fmt.Errorf("%w: repository has uncommitted changes; specify --base-ref or clean the repo", store.ErrInvalidState)
		}
	}

	baseCommit, err := resolveCommit(ctx, repoRoot, "HEAD")
	if err != nil {
		return "", "", protocol.InvalidInput("failed to resolve HEAD commit", err)
	}
	return "HEAD", baseCommit, nil
}

func resolveCommit(ctx context.Context, repoRoot, ref string) (string, error) {
	stdout, _, err := runGit(ctx, repoRoot, "rev-parse", "--verify", ref+"^{commit}")
	if err != nil {
		return "", err
	}
	return strings.TrimSpace(stdout), nil
}

func repoHasUncommittedChanges(ctx context.Context, repoRoot, workspaceRoot string) (bool, error) {
	stdout, _, err := runGit(ctx, repoRoot, "status", "--porcelain")
	if err != nil {
		return false, fmt.Errorf("check repository status: %w", err)
	}

	for _, line := range strings.Split(strings.TrimSpace(stdout), "\n") {
		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}
		if len(line) >= 3 {
			path := strings.TrimSpace(line[3:])
			if shouldIgnoreStatusPath(repoRoot, workspaceRoot, path) {
				continue
			}
		}
		return true, nil
	}

	return false, nil
}

func resolveWorkspaceRoot(repoRoot, configuredRoot string) string {
	root := strings.TrimSpace(configuredRoot)
	if root == "" {
		return filepath.Join(repoRoot, ".orch", "worktrees")
	}
	if filepath.IsAbs(root) {
		return root
	}
	return filepath.Join(repoRoot, root)
}

func ensureWorkspaceRootIgnored(repoRoot, workspaceRoot string) error {
	relative, err := filepath.Rel(repoRoot, workspaceRoot)
	if err != nil {
		return fmt.Errorf("resolve workspace root exclude path: %w", err)
	}
	if relative == "." || strings.HasPrefix(relative, "..") {
		return nil
	}

	pattern := filepath.ToSlash(relative)
	if !strings.HasSuffix(pattern, "/") {
		pattern += "/"
	}

	excludePath := filepath.Join(repoRoot, ".git", "info", "exclude")
	if err := os.MkdirAll(filepath.Dir(excludePath), 0o755); err != nil {
		return fmt.Errorf("create git info dir: %w", err)
	}

	content, err := os.ReadFile(excludePath)
	if err != nil && !os.IsNotExist(err) {
		return fmt.Errorf("read git exclude file: %w", err)
	}
	if strings.Contains(string(content), pattern) {
		return nil
	}

	appendContent := pattern + "\n"
	if len(content) > 0 && !strings.HasSuffix(string(content), "\n") {
		appendContent = "\n" + appendContent
	}
	file, err := os.OpenFile(excludePath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o644)
	if err != nil {
		return fmt.Errorf("open git exclude file: %w", err)
	}
	defer file.Close()

	if _, err := file.WriteString(appendContent); err != nil {
		return fmt.Errorf("append git exclude file: %w", err)
	}
	return nil
}

func shouldIgnoreStatusPath(repoRoot, workspaceRoot, statusPath string) bool {
	relative, err := filepath.Rel(repoRoot, workspaceRoot)
	if err != nil || relative == "." || strings.HasPrefix(relative, "..") {
		return false
	}

	relative = filepath.ToSlash(relative)
	statusPath = filepath.ToSlash(strings.Trim(statusPath, `"`))
	return statusPath == relative || strings.HasPrefix(statusPath, relative+"/")
}

func buildAttemptBranchName(runID, taskID string, attemptNo int) string {
	return fmt.Sprintf(
		"orch/%s/%s/attempt-%d",
		sanitizeGitSegment(runID),
		sanitizeGitSegment(taskID),
		attemptNo,
	)
}

func buildAttemptWorktreePath(workspaceRoot, runID, taskID string, attemptNo int) string {
	return filepath.Join(
		workspaceRoot,
		sanitizePathSegment(runID),
		sanitizePathSegment(taskID),
		fmt.Sprintf("attempt-%d", attemptNo),
	)
}

func deriveWorkspaceRootFromAttempt(runID, taskID, worktreePath string) (string, bool) {
	suffix := filepath.Join(
		sanitizePathSegment(runID),
		sanitizePathSegment(taskID),
		filepath.Base(worktreePath),
	)
	parent := filepath.Dir(worktreePath)
	if filepath.Base(parent) != sanitizePathSegment(taskID) {
		return "", false
	}
	runDir := filepath.Dir(parent)
	if filepath.Base(runDir) != sanitizePathSegment(runID) {
		return "", false
	}
	root := filepath.Dir(runDir)
	if filepath.Clean(filepath.Join(root, suffix)) != filepath.Clean(worktreePath) {
		return "", false
	}
	return root, true
}

func sanitizeGitSegment(value string) string {
	return sanitizeSegment(value)
}

func sanitizePathSegment(value string) string {
	return sanitizeSegment(value)
}

func sanitizeSegment(value string) string {
	value = strings.TrimSpace(value)
	if value == "" {
		return "item"
	}

	var b strings.Builder
	lastDash := false
	for _, r := range value {
		if (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') || (r >= '0' && r <= '9') {
			b.WriteRune(r)
			lastDash = false
			continue
		}
		if r == '-' || r == '_' || r == '.' {
			if !lastDash {
				b.WriteByte('-')
				lastDash = true
			}
			continue
		}
		if !lastDash {
			b.WriteByte('-')
			lastDash = true
		}
	}

	result := strings.Trim(b.String(), "-.")
	if result == "" {
		return "item"
	}
	if strings.HasSuffix(result, ".lock") {
		result = strings.TrimSuffix(result, ".lock") + "-lock"
	}
	return result
}

func runGit(ctx context.Context, repoRoot string, args ...string) (string, string, error) {
	cmdArgs := append([]string{"-C", repoRoot}, args...)
	cmd := exec.CommandContext(ctx, "git", cmdArgs...)
	output, err := cmd.CombinedOutput()
	if err == nil {
		return string(output), "", nil
	}

	message := strings.TrimSpace(string(output))
	if message == "" {
		message = err.Error()
	}
	return "", message, fmt.Errorf("git %s: %s", strings.Join(args, " "), message)
}

func cleanupAttemptWorktree(ctx context.Context, attempt store.TaskAttempt, force bool) error {
	if strings.TrimSpace(attempt.WorktreePath) == "" {
		return nil
	}

	if _, err := os.Stat(attempt.WorktreePath); err != nil {
		if os.IsNotExist(err) {
			return nil
		}
		return fmt.Errorf("stat worktree path: %w", err)
	}

	repoRoot, err := resolveRepoRootFromExistingWorktree(ctx, attempt.WorktreePath)
	if err != nil {
		if force {
			return os.RemoveAll(attempt.WorktreePath)
		}
		return err
	}

	args := []string{"worktree", "remove"}
	if force {
		args = append(args, "--force")
	}
	args = append(args, attempt.WorktreePath)
	if _, _, err := runGit(ctx, repoRoot, args...); err != nil {
		if force {
			return os.RemoveAll(attempt.WorktreePath)
		}
		return err
	}

	return nil
}

func resolveRepoRootFromExistingWorktree(ctx context.Context, worktreePath string) (string, error) {
	commonDir, err := resolveCommonGitDir(ctx, worktreePath)
	if err != nil {
		return "", err
	}
	return filepath.Dir(commonDir), nil
}

func resolveCommonGitDir(ctx context.Context, repoPath string) (string, error) {
	stdout, _, err := runGit(ctx, repoPath, "rev-parse", "--path-format=absolute", "--git-common-dir")
	if err != nil {
		return "", err
	}

	commonDir := strings.TrimSpace(stdout)
	if !filepath.IsAbs(commonDir) {
		commonDir = filepath.Join(repoPath, commonDir)
	}
	return filepath.Clean(commonDir), nil
}