package main

import (
	"fmt"
	"os"
	"regexp"
	"strings"
)

func addAuthorToFile(filePath string) error {
	content, err := os.ReadFile(filePath)
	if err != nil {
		return err
	}

	lines := strings.Split(string(content), "\n")
	var result []string
	var commentBuffer []string
	inComment := false

	funcPattern := regexp.MustCompile(`^func\s+(\([^)]+\)\s+)?\w+\s*\(`)

	for i := 0; i < len(lines); i++ {
		line := lines[i]
		trimmed := strings.TrimSpace(line)

		if strings.HasPrefix(trimmed, "//") && !strings.HasPrefix(trimmed, "//go:") {
			if strings.Contains(trimmed, "@author") {
				for _, c := range commentBuffer {
					result = append(result, c)
				}
				commentBuffer = nil
				result = append(result, line)
				inComment = false
				continue
			}
			commentBuffer = append(commentBuffer, line)
			inComment = true
			continue
		}

		if inComment && len(commentBuffer) > 0 && funcPattern.MatchString(trimmed) {
			for _, c := range commentBuffer {
				result = append(result, c)
			}
			indent := ""
			for _, ch := range line {
				if ch == ' ' || ch == '\t' {
					indent += string(ch)
				} else {
					break
				}
			}
			result = append(result, indent+"// @author sunct")
			result = append(result, line)
			commentBuffer = nil
			inComment = false
			continue
		}

		if len(commentBuffer) > 0 {
			for _, c := range commentBuffer {
				result = append(result, c)
			}
			commentBuffer = nil
		}
		inComment = false
		result = append(result, line)
	}

	if len(commentBuffer) > 0 {
		for _, c := range commentBuffer {
			result = append(result, c)
		}
	}

	output := strings.Join(result, "\n")
	return os.WriteFile(filePath, []byte(output), 0644)
}

func main() {
	if len(os.Args) < 2 {
		fmt.Println("Usage: go run add_author.go <file1.go> [file2.go ...]")
		os.Exit(1)
	}

	for _, file := range os.Args[1:] {
		fmt.Printf("Processing: %s ... ", file)
		err := addAuthorToFile(file)
		if err != nil {
			fmt.Printf("ERROR: %v\n", err)
		} else {
			fmt.Println("OK")
		}
	}
}