CodeCommentRemover.go

//go:build go1.25

package main

import (
	"bytes"
	"cmp"
	"errors"
	"flag"
	"fmt"
	"io/fs"
	"maps"
	"os"
	"path/filepath"
	"runtime"
	"slices"
	"strings"
	"sync"
	"unicode"
)

type blockPattern struct {
	start string
	end   string
}

type commentPatterns struct {
	linePrefixes []string
	blockTokens  []blockPattern
}

func cloneCommentPatterns(src commentPatterns) commentPatterns {
	return commentPatterns{
		linePrefixes: slices.Clone(src.linePrefixes),
		blockTokens:  slices.Clone(src.blockTokens),
	}
}

func mergeCommentPatterns(dst *commentPatterns, addition commentPatterns) {
	dst.linePrefixes = append(dst.linePrefixes, addition.linePrefixes...)
	dst.blockTokens = append(dst.blockTokens, addition.blockTokens...)
}

type quoteState struct {
	inSingle   bool
	inDouble   bool
	inBacktick bool
	escaped    bool
}

func (s *quoteState) reset() {
	*s = quoteState{}
}

func (s *quoteState) clearTransient() {
	s.escaped = false
	s.inSingle = false
	s.inDouble = false
}

func (s *quoteState) consumeEscaped() bool {
	if s.escaped {
		s.escaped = false
		return true
	}
	return false
}

func (s *quoteState) advance(ch byte) {
	if s.escaped {
		s.escaped = false
		return
	}

	switch ch {
	case '\\':
		if s.inSingle || s.inDouble || s.inBacktick {
			s.escaped = true
		}
	case '`':
		if !s.inSingle && !s.inDouble {
			s.inBacktick = !s.inBacktick
		}
	case '\'':
		if !s.inDouble && !s.inBacktick {
			s.inSingle = !s.inSingle
		}
	case '"':
		if !s.inSingle && !s.inBacktick {
			s.inDouble = !s.inDouble
		}
	}
}

func (s *quoteState) inQuote() bool {
	return s.inSingle || s.inDouble || s.inBacktick
}

var aliasPatterns = map[string]string{
	"standard":   "//",
	"block":      "/*|*/",
	"doxygen":    "//,/**|*/,/*|*/",
	"cpp":        "//,/*|*/",
	"shell":      "#",
	"hash":       "#",
	"go":         "//,/*|*/",
	"c":          "//,/*|*/",
	"csharp":     "//,/*|*/",
	"typescript": "//,/*|*/",
}

var compiledAliasPatterns map[string]commentPatterns

func init() {
	compiledAliasPatterns = make(map[string]commentPatterns, len(aliasPatterns))
	for name, spec := range aliasPatterns {
		parsed, err := parsePatternInternal(spec, nil)
		if err != nil {
			panic(fmt.Sprintf("invalid alias %q specification: %v", name, err))
		}
		compiledAliasPatterns[name] = parsed
	}
}

func patternsForAlias(name string) (commentPatterns, error) {
	aliasName := strings.ToLower(strings.TrimSpace(name))
	patterns, ok := compiledAliasPatterns[aliasName]
	if !ok {
		return commentPatterns{}, fmt.Errorf("unknown alias %q", name)
	}
	return cloneCommentPatterns(patterns), nil
}

func main() {
	workers := flag.Int("workers", 0, "max files to process concurrently (0 uses GOMAXPROCS)")
	flag.Usage = func() {
		fmt.Fprintf(flag.CommandLine.Output(), "Usage: %s [comment pattern] [file|folder]\n", os.Args[0])
		fmt.Fprintf(flag.CommandLine.Output(), "   or: %s AliasMode <alias> [file|folder]\n", os.Args[0])
		fmt.Fprintln(flag.CommandLine.Output(), "\nPattern examples:")
		fmt.Fprintln(flag.CommandLine.Output(), "  //             remove lines that start with //")
		fmt.Fprintln(flag.CommandLine.Output(), "  //,#           remove lines starting with // or #")
		fmt.Fprintln(flag.CommandLine.Output(), "  \"/*|*/\"        remove /* ... */ blocks (quote on Windows!)")
		fmt.Fprintln(flag.CommandLine.Output(), "  \"//,/*|*/\"     combine line and block definitions")
		fmt.Fprintln(flag.CommandLine.Output(), "\nAliases:")
		names := slices.Collect(maps.Keys(aliasPatterns))
		slices.Sort(names)
		for _, name := range names {
			fmt.Fprintf(flag.CommandLine.Output(), "  %-11s -> %s\n", name, aliasPatterns[name])
		}
		fmt.Fprintln(flag.CommandLine.Output(), "\nFlags:")
		flag.PrintDefaults()
	}

	flag.Parse()

	desiredProcs := cmp.Or(runtime.NumCPU(), 1)
	currentProcs := runtime.GOMAXPROCS(0)
	if currentProcs < desiredProcs {
		runtime.GOMAXPROCS(desiredProcs)
		currentProcs = desiredProcs
	}

	args := flag.Args()
	var (
		targetPath  string
		patternSpec commentPatterns
		err         error
	)

	if len(args) > 0 && strings.EqualFold(args[0], "AliasMode") {
		if len(args) < 2 || len(args) > 3 {
			flag.Usage()
			os.Exit(1)
		}
		aliasName := args[1]
		targetPath = "."
		if len(args) == 3 {
			targetPath = args[2]
		}
		patternSpec, err = patternsForAlias(aliasName)
	} else {
		if len(args) != 2 {
			flag.Usage()
			os.Exit(1)
		}
		patternSpec, err = parsePattern(args[0])
		targetPath = args[1]
	}

	if err != nil {
		fmt.Fprintf(os.Stderr, "Error parsing pattern: %v\n", err)
		os.Exit(1)
	}

	info, err := os.Stat(targetPath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error accessing %s: %v\n", targetPath, err)
		os.Exit(1)
	}

	workerLimit := cmp.Or(*workers, currentProcs)
	if workerLimit < 1 {
		workerLimit = 1
	}
	if workerLimit > currentProcs {
		runtime.GOMAXPROCS(workerLimit)
		currentProcs = workerLimit
	}

	var (
		processErr     error
		processedCount int
		skippedCount   int
		countMu        sync.Mutex
	)

	if info.IsDir() {
		var (
			joined   error
			joinedMu sync.Mutex
			wg       sync.WaitGroup
		)

		sem := make(chan struct{}, workerLimit)

		processErr = filepath.WalkDir(targetPath, func(path string, d fs.DirEntry, walkErr error) error {
			if walkErr != nil {
				return walkErr
			}

			if d.IsDir() {
				return nil
			}

			sem <- struct{}{}
			wg.Go(func() {
				defer func() { <-sem }()

				modified, err := processFile(path, patternSpec)
				if err != nil {
					fmt.Fprintf(os.Stderr, "Failed %s: %v\n", path, err)
					joinedMu.Lock()
					joined = errors.Join(joined, err)
					joinedMu.Unlock()
					return
				}
				countMu.Lock()
				if modified {
					processedCount++
				} else {
					skippedCount++
				}
				countMu.Unlock()
			})
			return nil
		})
		wg.Wait()
		if processErr == nil {
			processErr = joined
		} else if joined != nil {
			processErr = errors.Join(processErr, joined)
		}
	} else {
		modified, err := processFile(targetPath, patternSpec)
		processErr = err
		if modified {
			processedCount = 1
		} else {
			skippedCount = 1
		}
	}

	if processErr != nil {
		fmt.Fprintf(os.Stderr, "Processing error: %v\n", processErr)
		os.Exit(1)
	}

	if processedCount == 0 && skippedCount == 0 {
		fmt.Println("No files found to process.")
	} else if processedCount == 0 {
		fmt.Printf("No changes needed. Checked %d file(s), no comments matched the pattern.\n", skippedCount)
	} else {
		fmt.Printf("Completed: %d file(s) updated, %d file(s) unchanged.\n", processedCount, skippedCount)
	}
}

func parsePattern(raw string) (commentPatterns, error) {
	return parsePatternInternal(raw, func(name string) (commentPatterns, bool) {
		addition, ok := compiledAliasPatterns[name]
		return addition, ok
	})
}

func parsePatternInternal(raw string, aliasResolver func(string) (commentPatterns, bool)) (commentPatterns, error) {
	raw = strings.TrimSpace(raw)
	if raw == "" {
		return commentPatterns{}, errors.New("empty comment pattern")
	}

	patterns := commentPatterns{}

	for token := range strings.SplitSeq(raw, ",") {
		token = strings.TrimSpace(token)
		if token == "" {
			continue
		}

		if aliasResolver != nil {
			if addition, ok := aliasResolver(strings.ToLower(token)); ok {
				mergeCommentPatterns(&patterns, addition)
				continue
			}
		}

		if start, end, found := strings.Cut(token, "|"); found {
			start = strings.TrimSpace(start)
			end = strings.TrimSpace(end)
			if start == "" || end == "" {
				return commentPatterns{}, fmt.Errorf("invalid block pattern: %q", token)
			}
			patterns.blockTokens = append(patterns.blockTokens, blockPattern{start: start, end: end})
			continue
		}

		patterns.linePrefixes = append(patterns.linePrefixes, token)
	}

	if len(patterns.linePrefixes) > 1 {
		slices.Sort(patterns.linePrefixes)
		patterns.linePrefixes = slices.Compact(patterns.linePrefixes)
		patterns.linePrefixes = slices.Clip(patterns.linePrefixes)
	}

	if len(patterns.blockTokens) > 1 {
		slices.SortFunc(patterns.blockTokens, func(a, b blockPattern) int {
			if diff := cmp.Compare(a.start, b.start); diff != 0 {
				return diff
			}
			return cmp.Compare(a.end, b.end)
		})
		patterns.blockTokens = slices.CompactFunc(patterns.blockTokens, func(prev, next blockPattern) bool {
			return prev.start == next.start && prev.end == next.end
		})
		patterns.blockTokens = slices.Clip(patterns.blockTokens)
	}

	if len(patterns.linePrefixes) == 0 && len(patterns.blockTokens) == 0 {
		return commentPatterns{}, errors.New("no valid patterns found")
	}

	return patterns, nil
}

func processFile(path string, patterns commentPatterns) (bool, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return false, fmt.Errorf("read %s: %w", path, err)
	}

	if isLikelyBinary(data) {
		return false, nil
	}

	original := string(data)
	cleaned := stripComments(original, patterns)

	if cleaned == original {
		return false, nil
	}

	info, err := os.Stat(path)
	if err != nil {
		return false, fmt.Errorf("stat %s: %w", path, err)
	}

	if err := os.WriteFile(path, []byte(cleaned), info.Mode().Perm()); err != nil {
		return false, fmt.Errorf("write %s: %w", path, err)
	}

	fmt.Printf("Updated %s\n", path)
	return true, nil
}

func stripComments(content string, patterns commentPatterns) string {
	// Process block comments first, then line comments
	// removeBlockComments is now aware of line comment prefixes to avoid quote state corruption
	withoutBlocks := content
	for _, block := range patterns.blockTokens {
		withoutBlocks = removeBlockComments(withoutBlocks, block, patterns.linePrefixes)
	}

	if len(patterns.linePrefixes) > 0 {
		return removeLineComments(withoutBlocks, patterns.linePrefixes)
	}

	return withoutBlocks
}

func removeBlockComments(content string, block blockPattern, linePrefixes []string) string {
	if block.start == "" || block.end == "" {
		return content
	}

	var builder strings.Builder
	builder.Grow(len(content))
	state := quoteState{}

	i := 0
	for i < len(content) {
		if state.consumeEscaped() {
			builder.WriteByte(content[i])
			i++
			continue
		}

		// Check if we're at the start of a line comment - if so, skip the entire line
		// This prevents quote characters in line comments from corrupting our quote state
		isLineComment := false
		if !state.inQuote() && len(linePrefixes) > 0 {
			for _, prefix := range linePrefixes {
				if strings.HasPrefix(content[i:], prefix) {
					// Found a line comment - copy it to output and skip to next line
					lineEnd := strings.IndexAny(content[i:], "\r\n")
					if lineEnd == -1 {
						// Comment goes to end of file
						builder.WriteString(content[i:])
						i = len(content)
					} else {
						builder.WriteString(content[i : i+lineEnd])
						i += lineEnd
					}
					isLineComment = true
					break
				}
			}
		}
		if isLineComment {
			continue
		}

		if !state.inQuote() && strings.HasPrefix(content[i:], block.start) {
			endIdx := strings.Index(content[i+len(block.start):], block.end)
			if endIdx == -1 {
				// Unclosed block comment - skip everything from here to end of file
				break
			}
			// Properly closed comment - skip the entire block
			i += len(block.start) + endIdx + len(block.end)
			state.reset()
			continue
		}

		ch := content[i]
		state.advance(ch)

		builder.WriteByte(ch)
		i++
	}

	return builder.String()
}

func removeLineComments(content string, prefixes []string) string {
	var builder strings.Builder
	builder.Grow(len(content))
	state := quoteState{}
	lineIndex := 0

	for offset := 0; offset < len(content); {
		line, newline, next := scanLinePreservingEnding(content, offset)

		if isShebangLine(line, lineIndex) {
			state.reset()
			builder.WriteString(line)
			builder.WriteString(newline)
			offset = next
			lineIndex++
			continue
		}
		nonSpaceIdx := indexFirstNonSpace(line)

		if idx := findFirstComment(line, prefixes, &state); idx >= 0 {
			if nonSpaceIdx >= 0 && idx == nonSpaceIdx {
				state.clearTransient()
				offset = next
				lineIndex++
				continue
			}
			line = line[:idx]
			// Trim trailing whitespace but preserve CR/LF characters
			line = trimTrailingWhitespacePreserveCR(line)
			// If line is now empty or only whitespace after removing comment, skip it entirely
			if len(line) == 0 || indexFirstNonSpace(line) == -1 {
				state.clearTransient()
				offset = next
				lineIndex++
				continue
			}
		}

		state.clearTransient()

		builder.WriteString(line)
		builder.WriteString(newline)
		offset = next
		lineIndex++
	}

	return builder.String()
}

func findFirstComment(line string, prefixes []string, state *quoteState) int {
	if len(prefixes) == 0 {
		return -1
	}

	for i := range len(line) {
		ch := line[i]

		if state.consumeEscaped() {
			continue
		}

		state.advance(ch)

		if state.inQuote() {
			continue
		}

		for _, prefix := range prefixes {
			if strings.HasPrefix(line[i:], prefix) {
				// Preserve "// namespace" comments (C++ namespace closing brace annotations)
				if isNamespaceComment(line, i) {
					continue
				}
				return i
			}
		}
	}

	return -1
}

func isNamespaceComment(line string, commentStart int) bool {
	// Validate index bounds
	if commentStart < 0 || commentStart >= len(line) {
		return false
	}

	segment := line[commentStart:]

	// Check if this is a "//" comment and extract content after it
	afterComment, found := strings.CutPrefix(segment, "//")
	if !found {
		return false
	}

	// Extract and normalize the text after "//"
	rest := strings.TrimSpace(afterComment)
	restLower := strings.ToLower(rest)

	// Check if it starts with "namespace" as a complete word
	afterNamespace, found := strings.CutPrefix(restLower, "namespace")
	if !found {
		return false
	}

	// If nothing after "namespace", it's a valid namespace comment
	if afterNamespace == "" {
		return true
	}

	// Ensure "namespace" is followed by whitespace (not part of a longer word like "namespaced")
	return unicode.IsSpace(rune(afterNamespace[0]))
}

func scanLinePreservingEnding(content string, start int) (string, string, int) {
	if start >= len(content) {
		return "", "", len(content)
	}

	segment := content[start:]
	if idx := strings.IndexAny(segment, "\r\n"); idx >= 0 {
		newlineLen := 1
		if segment[idx] == '\r' && idx+1 < len(segment) && segment[idx+1] == '\n' {
			newlineLen = 2
		}
		return segment[:idx], segment[idx : idx+newlineLen], start + idx + newlineLen
	}

	return segment, "", len(content)
}

func indexFirstNonSpace(s string) int {
	if idx := strings.IndexFunc(s, func(r rune) bool {
		return !unicode.IsSpace(r)
	}); idx >= 0 {
		return idx
	}
	return -1
}

func trimTrailingWhitespacePreserveCR(s string) string {
	return strings.TrimRightFunc(s, func(r rune) bool {
		return unicode.IsSpace(r) && r != '\r' && r != '\n'
	})
}

func isShebangLine(line string, index int) bool {
	if index != 0 {
		return false
	}
	if trimmed, ok := strings.CutPrefix(line, "\ufeff"); ok {
		line = trimmed
	}
	return strings.HasPrefix(line, "#!")
}

func isLikelyBinary(data []byte) bool {
	return bytes.IndexByte(data, 0) >= 0
}