onyx-prebootstrap/internal/storage/state.go

package storage

import (
	"fmt"

	"git.dws.rip/DWS/onyx/internal/models"
	"github.com/go-git/go-git/v5"
	"github.com/go-git/go-git/v5/plumbing"
)

// StateCapture provides functionality to capture repository state
type StateCapture struct {
	repo *git.Repository
}

// NewStateCapture creates a new StateCapture instance
func NewStateCapture(repo *git.Repository) *StateCapture {
	return &StateCapture{
		repo: repo,
	}
}

// CaptureState captures the current state of the repository
func (s *StateCapture) CaptureState() (*models.RepositoryState, error) {
	refs, err := s.captureRefs()
	if err != nil {
		return nil, fmt.Errorf("failed to capture refs: %w", err)
	}

	currentWorkstream, err := s.getCurrentWorkstream()
	if err != nil {
		// It's okay if there's no current workstream (e.g., in detached HEAD state)
		currentWorkstream = ""
	}

	workingTreeHash, err := s.getWorkingTreeHash()
	if err != nil {
		// Working tree hash might not be available in a fresh repo
		workingTreeHash = ""
	}

	indexHash, err := s.getIndexHash()
	if err != nil {
		// Index hash might not be available in a fresh repo
		indexHash = ""
	}

	return models.NewRepositoryState(refs, currentWorkstream, workingTreeHash, indexHash), nil
}

// captureRefs captures all Git references (branches, tags, etc.)
func (s *StateCapture) captureRefs() (map[string]string, error) {
	refs := make(map[string]string)

	refIter, err := s.repo.References()
	if err != nil {
		return nil, fmt.Errorf("failed to get references: %w", err)
	}

	err = refIter.ForEach(func(ref *plumbing.Reference) error {
		if ref.Type() == plumbing.HashReference {
			refs[ref.Name().String()] = ref.Hash().String()
		} else if ref.Type() == plumbing.SymbolicReference {
			// For symbolic refs (like HEAD), store the target
			refs[ref.Name().String()] = ref.Target().String()
		}
		return nil
	})

	if err != nil {
		return nil, fmt.Errorf("failed to iterate references: %w", err)
	}

	return refs, nil
}

// getCurrentWorkstream determines the current workstream (branch)
func (s *StateCapture) getCurrentWorkstream() (string, error) {
	head, err := s.repo.Head()
	if err != nil {
		return "", fmt.Errorf("failed to get HEAD: %w", err)
	}

	if head.Name().IsBranch() {
		return head.Name().Short(), nil
	}

	// In detached HEAD state
	return "", fmt.Errorf("in detached HEAD state")
}

// getWorkingTreeHash gets a hash representing the current working tree
func (s *StateCapture) getWorkingTreeHash() (string, error) {
	worktree, err := s.repo.Worktree()
	if err != nil {
		return "", fmt.Errorf("failed to get worktree: %w", err)
	}

	status, err := worktree.Status()
	if err != nil {
		return "", fmt.Errorf("failed to get status: %w", err)
	}

	// For now, we'll just check if the working tree is clean
	// In the future, we might compute an actual hash
	if status.IsClean() {
		head, err := s.repo.Head()
		if err == nil {
			return head.Hash().String(), nil
		}
	}

	return "dirty", nil
}

// getIndexHash gets a hash representing the current index (staging area)
func (s *StateCapture) getIndexHash() (string, error) {
	// For now, this is a placeholder
	// In the future, we might compute a proper hash of the index
	return "", nil
}

// RestoreState restores the repository to a previously captured state
func (s *StateCapture) RestoreState(state *models.RepositoryState) error {
	// Restore all refs
	for refName, refHash := range state.Refs {
		ref := plumbing.NewReferenceFromStrings(refName, refHash)

		// Skip symbolic references for now
		if ref.Type() == plumbing.SymbolicReference {
			continue
		}

		err := s.repo.Storer.SetReference(ref)
		if err != nil {
			return fmt.Errorf("failed to restore ref %s: %w", refName, err)
		}
	}

	// If there's a current workstream, check it out
	if state.CurrentWorkstream != "" {
		worktree, err := s.repo.Worktree()
		if err != nil {
			return fmt.Errorf("failed to get worktree: %w", err)
		}

		err = worktree.Checkout(&git.CheckoutOptions{
			Branch: plumbing.NewBranchReferenceName(state.CurrentWorkstream),
		})
		if err != nil {
			// Don't fail if checkout fails, just log it
			// The refs have been restored which is the most important part
			fmt.Printf("Warning: failed to checkout branch %s: %v\n", state.CurrentWorkstream, err)
		}
	}

	return nil
}

// CompareStates compares two repository states and returns the differences
func (s *StateCapture) CompareStates(before, after *models.RepositoryState) map[string]string {
	differences := make(map[string]string)

	// Check for changed/added refs
	for refName, afterHash := range after.Refs {
		beforeHash, exists := before.Refs[refName]
		if !exists {
			differences[refName] = fmt.Sprintf("added: %s", afterHash)
		} else if beforeHash != afterHash {
			differences[refName] = fmt.Sprintf("changed: %s -> %s", beforeHash, afterHash)
		}
	}

	// Check for deleted refs
	for refName := range before.Refs {
		if _, exists := after.Refs[refName]; !exists {
			differences[refName] = "deleted"
		}
	}

	// Check workstream change
	if before.CurrentWorkstream != after.CurrentWorkstream {
		differences["current_workstream"] = fmt.Sprintf("changed: %s -> %s", before.CurrentWorkstream, after.CurrentWorkstream)
	}

	return differences
}