internal/rewriter/rewriter.go

Ref: main   Size: 3.2 KiB

package rewriter

import (
	"bufio"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"time"
)

var ErrConcurrentModification = errors.New("file modified concurrently")

// stampRE matches "@claude" optionally followed by an existing "[fw-XXXXXXXX]" tag.
var stampRE = regexp.MustCompile(`@claude(?:\[fw-[0-9a-f]{8}\])?`)

// StampUUIDs rewrites the file at path, splicing "[fw-XXXXXXXX]" immediately
// after the first untagged "@claude" on each given line. stamps maps line
// number (1-indexed) → UUID (e.g. "fw-a1b2c3d4"). The write is atomic via
// tmpfile + rename. File mode is preserved.
//
// Aborts with ErrConcurrentModification if the file's mtime changed between
// read and rename (someone — likely the user's editor — wrote concurrently).
// In that case, the rewrite is dropped and the marker stays untagged for a
// retry on the next save event.
func StampUUIDs(path string, stamps map[int]string) error {
	info, err := os.Stat(path)
	if err != nil {
		return err
	}
	mtimeBefore := info.ModTime()

	in, err := os.Open(path)
	if err != nil {
		return err
	}
	defer in.Close()

	dir := filepath.Dir(path)
	tmp, err := os.CreateTemp(dir, ".filewatch-*.tmp")
	if err != nil {
		return err
	}
	tmpPath := tmp.Name()
	cleanup := func() {
		tmp.Close()
		os.Remove(tmpPath)
	}

	sc := bufio.NewScanner(in)
	sc.Buffer(make([]byte, 1024*1024), 1024*1024)
	bw := bufio.NewWriter(tmp)
	line := 0
	for sc.Scan() {
		line++
		text := sc.Text()
		if uuid, ok := stamps[line]; ok {
			text = stampLine(text, uuid)
		}
		if _, err := bw.WriteString(text); err != nil {
			cleanup()
			return err
		}
		if _, err := bw.WriteString("\n"); err != nil {
			cleanup()
			return err
		}
	}
	if err := sc.Err(); err != nil {
		cleanup()
		return err
	}
	if err := bw.Flush(); err != nil {
		cleanup()
		return err
	}
	if err := tmp.Chmod(info.Mode().Perm()); err != nil {
		cleanup()
		return err
	}
	if err := tmp.Close(); err != nil {
		os.Remove(tmpPath)
		return err
	}

	return atomicReplace(path, tmpPath, info.Mode().Perm(), mtimeBefore)
}

// atomicReplace stat-checks path against mtimeBefore. If the file was modified
// concurrently, returns ErrConcurrentModification and removes tmpPath. Otherwise
// renames tmpPath over path. The mode argument is currently unused (the caller
// already chmods tmpPath before close), but is included to keep this helper's
// signature self-contained for direct testing.
func atomicReplace(path, tmpPath string, mode os.FileMode, mtimeBefore time.Time) error {
	info, err := os.Stat(path)
	if err != nil {
		os.Remove(tmpPath)
		return err
	}
	if !info.ModTime().Equal(mtimeBefore) {
		os.Remove(tmpPath)
		return ErrConcurrentModification
	}
	if err := os.Rename(tmpPath, path); err != nil {
		os.Remove(tmpPath)
		return fmt.Errorf("rename: %w", err)
	}
	_ = mode // currently unused; reserved for future modes that need post-rename Chmod
	return nil
}

// stampLine splices "[uuid]" after the first untagged "@claude" on the line.
// Already-tagged occurrences are left untouched.
func stampLine(text, uuid string) string {
	done := false
	return stampRE.ReplaceAllStringFunc(text, func(m string) string {
		if done || strings.Contains(m, "[fw-") {
			return m
		}
		done = true
		return "@claude[" + uuid + "]"
	})
}