dirmd/internal/renderer/renderer.go

package renderer

import (
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"
	"unicode/utf8"

	"git.flo.fo/FLO/dirmd/internal/config"
	"git.flo.fo/FLO/dirmd/internal/filter"
	"git.flo.fo/FLO/dirmd/internal/git"
	"git.flo.fo/FLO/dirmd/internal/instructions"
	"git.flo.fo/FLO/dirmd/internal/tree"
	"git.flo.fo/FLO/dirmd/internal/walker"
)

const fenceChars = "~~~"

func Write(result *walker.Result, outputPath string, appendMode bool, singleFile bool) error {
	flags := os.O_CREATE | os.O_WRONLY
	if appendMode {
		flags |= os.O_APPEND
	} else {
		flags |= os.O_TRUNC
	}

	f, err := os.OpenFile(outputPath, flags, 0644)
	if err != nil {
		return fmt.Errorf("cannot open output file: %w", err)
	}
	defer f.Close()

	var sb strings.Builder

	h1Path := result.AbsRoot
	if singleFile && len(result.Entries) > 0 {
		h1Path = filepath.Join(result.AbsRoot, result.Entries[0].RelPath)
	}

	fmt.Fprintf(&sb, "# %s [ROOT]\n", h1Path)

	if !appendMode {
		fmt.Fprintf(&sb, "> Generated: %s\n", time.Now().Format(time.RFC3339))
		sb.WriteString(git.FormatCommitLine(result.Commit))

		if result.Instructions != "" {
			sb.WriteString(instructions.Format(result.Instructions))
		} else {
			sb.WriteString("\n")
		}

		sb.WriteString(FormatFlags(result.Cfg))
	}

	if !singleFile {
		treeRoot, err := tree.BuildFullTree(result.AbsRoot, result.Cfg)
		if err != nil {
			return fmt.Errorf("cannot build tree: %w", err)
		}

		rootName := filepath.Base(result.AbsRoot)
		sb.WriteString("## Directory Tree\n\n")
		sb.WriteString(fenceChars + "\n")
		sb.WriteString(renderTreeWithLegend(treeRoot, rootName))
		sb.WriteString(fenceChars + "\n\n")
	}

	if len(result.Entries) > 0 {
		if !singleFile {
			sb.WriteString("## Contents\n\n")
		}

		for _, entry := range result.Entries {
			fullPath := filepath.Join(result.AbsRoot, entry.RelPath)
			content, err := readFileForOutput(fullPath, result.OmitTest)
			if err != nil {
				fmt.Fprintf(os.Stderr, "warning: cannot read %s: %v\n", entry.RelPath, err)
				continue
			}

			language := extToLang(filepath.Ext(entry.RelPath))
			lines := countLines(content)

			fmt.Fprintf(&sb, "--- FILE: ./%s (%d bytes, %d lines) ---\n", entry.RelPath, len(content), lines)
			fmt.Fprintf(&sb, "%s", fenceChars)
			if language != "" {
				fmt.Fprintf(&sb, "%s", language)
			}
			fmt.Fprintf(&sb, "\n")
			fmt.Fprintf(&sb, "%s", content)
			if !strings.HasSuffix(content, "\n") {
				fmt.Fprintf(&sb, "\n")
			}
			fmt.Fprintf(&sb, "%s\n--- END FILE ---\n\n", fenceChars)
		}
	}

	_, err = f.WriteString(sb.String())
	return err
}

func WriteSlices(result *walker.Result, outputDir, repoName string) ([]string, error) {
	entries := result.Entries
	if len(entries) == 0 {
		return nil, fmt.Errorf("no entries to write")
	}

	absRoot := result.AbsRoot
	rootName := filepath.Base(absRoot)

	topLevelDirs := make(map[string][]walker.Entry)
	rootFiles := make([]walker.Entry, 0)

	for _, e := range entries {
		rel := e.RelPath
		parts := strings.Split(rel, string(filepath.Separator))
		if len(parts) == 1 {
			rootFiles = append(rootFiles, e)
		} else {
			firstPart := parts[0]
			topLevelDirs[firstPart] = append(topLevelDirs[firstPart], e)
		}
	}

	treeRoot, err := tree.BuildFullTree(absRoot, result.Cfg)
	if err != nil {
		return nil, fmt.Errorf("cannot build full tree: %w", err)
	}

	var dirNames []string
	for k := range topLevelDirs {
		dirNames = append(dirNames, k)
	}
	sort.Strings(dirNames)

	for _, dirName := range dirNames {
		if node, ok := treeRoot.Children[dirName]; ok {
			node.Ref = filepath.Join(dirName, dirName+"_"+repoName+".md")
		}
	}

	var generated []string
	indexBasename := "index_" + repoName + ".md"
	rootPath := filepath.Join(outputDir, indexBasename)
	generated = append(generated, rootPath)

	rootSliceBasename := "root_" + repoName + ".md"
	rootSlicePath := filepath.Join(outputDir, rootSliceBasename)
	generated = append(generated, rootSlicePath)

	pickInfos := make([]PickInfo, 0)

	if len(result.Cfg.PickPaths) > 0 {
		pickInfos, err = writePickFiles(absRoot, outputDir, result.Cfg, result.OmitTest, result.Commit)
		if err != nil {
			return nil, fmt.Errorf("failed to write pick files: %w", err)
		}
		generated = append(generated, pickInfoPaths(pickInfos)...)
	}

	if err := writeRootIndex(rootPath, absRoot, rootName, repoName, topLevelDirs, rootFiles, result.Commit, result.Instructions, treeRoot, result.Cfg, pickInfos); err != nil {
		return nil, fmt.Errorf("failed to write root index: %w", err)
	}

	if err := writeRootSliceFile(rootSlicePath, absRoot, rootFiles, result.OmitTest, result.Commit, result.Cfg); err != nil {
		return nil, fmt.Errorf("failed to write root slice: %w", err)
	}

	for _, dirName := range dirNames {
		sliceName := dirName + "_" + repoName + ".md"
		sliceDir := filepath.Join(outputDir, dirName)
		if err := os.MkdirAll(sliceDir, 0755); err != nil {
			return nil, fmt.Errorf("cannot create slice directory %s: %w", sliceDir, err)
		}
		slicePath := filepath.Join(sliceDir, sliceName)
		generated = append(generated, slicePath)

		absDir := filepath.Join(absRoot, dirName)
		dirEntries := topLevelDirs[dirName]

		if err := writeDirSlice(slicePath, absDir, absRoot, dirName, dirEntries, indexBasename, result.OmitTest, result.Commit, result.Cfg); err != nil {
			return nil, fmt.Errorf("failed to write slice %s: %w", sliceName, err)
		}
	}

	return generated, nil
}

type PickInfo struct {
	RelPath  string
	FileType string
	SliceDir string
}

func pickInfoPaths(picks []PickInfo) []string {
	var paths []string
	for _, p := range picks {
		paths = append(paths, filepath.Join(p.SliceDir, p.FileType))
	}
	return paths
}

func writePickFiles(absRoot, outputDir string, cfg *config.Config, omitTest bool, commit *git.CommitInfo) ([]PickInfo, error) {
	var pickInfos []PickInfo

	for _, pickPath := range cfg.PickPaths {
		fullPath := filepath.Join(absRoot, pickPath)

		info, err := os.Stat(fullPath)
		if err != nil {
			return nil, fmt.Errorf("cannot access pick path %s: %w", pickPath, err)
		}

		if info.IsDir() {
			return nil, fmt.Errorf("pick path must be a file: %s", pickPath)
		}

		name := filepath.Base(pickPath)
		flt := &filter.Filter{
			MaxSize:      cfg.MaxSize,
			Extensions:   cfg.Extensions,
			Ignores:      cfg.Ignores,
			OmitTest:     cfg.OmitTest,
			FrontendOnly: cfg.FrontendOnly,
			OmitMd:       cfg.OmitMd,
		}

		if flt.ShouldSkipFile(fullPath, name, info) {
			return nil, fmt.Errorf("pick path %s filtered by active rules", pickPath)
		}

		parts := strings.Split(pickPath, string(filepath.Separator))
		sliceDir := outputDir
		if len(parts) > 1 {
			sliceDir = filepath.Join(outputDir, parts[0])
			if err := os.MkdirAll(sliceDir, 0755); err != nil {
				return nil, fmt.Errorf("cannot create pick slice directory %s: %w", sliceDir, err)
			}
		}

		sliceName := filepath.Base(pickPath) + ".md"
		pickPathOut := filepath.Join(sliceDir, sliceName)

		if err := writePickFile(pickPathOut, absRoot, pickPath, omitTest, commit, cfg); err != nil {
			return nil, fmt.Errorf("failed to write pick %s: %w", pickPath, err)
		}

		pickInfos = append(pickInfos, PickInfo{
			RelPath:  pickPath,
			FileType: sliceName,
			SliceDir: filepath.Base(sliceDir),
		})
	}

	return pickInfos, nil
}

func writePickFile(path, absRoot, relPath string, omitTest bool, commit *git.CommitInfo, cfg *config.Config) error {
	f, err := os.Create(path)
	if err != nil {
		return fmt.Errorf("cannot create pick file: %w", err)
	}
	defer f.Close()

	var sb strings.Builder

	fmt.Fprintf(&sb, "# %s [SLICE]\n", filepath.Join(absRoot, relPath))
	fmt.Fprintf(&sb, "> Generated: %s\n", time.Now().Format(time.RFC3339))
	sb.WriteString(git.FormatCommitLine(commit))
	sb.WriteString(FormatFlags(cfg))
	fmt.Fprintf(&sb, "Root document: index_%s.md\n\n", filepath.Base(absRoot))

	fullPath := filepath.Join(absRoot, relPath)
	content, err := readFileForOutput(fullPath, omitTest)
	if err != nil {
		return fmt.Errorf("cannot read pick file %s: %w", relPath, err)
	}
	language := extToLang(filepath.Ext(relPath))
	lines := countLines(content)

	fmt.Fprintf(&sb, "Tree: %s\n\n", tree.RenderBreadcrumb(absRoot, relPath))

	fmt.Fprintf(&sb, "--- FILE: ./%s (%d bytes, %d lines) ---\n", relPath, len(content), lines)
	sb.WriteString(fenceChars)
	if language != "" {
		sb.WriteString(language)
	}
	sb.WriteString("\n")
	sb.WriteString(content)
	if !strings.HasSuffix(content, "\n") {
		sb.WriteString("\n")
	}
	fmt.Fprintf(&sb, "%s\n--- END FILE ---\n\n", fenceChars)

	_, err = f.WriteString(sb.String())
	return err
}

func writeRootIndex(path, absRoot, rootName, repoName string, topLevelDirs map[string][]walker.Entry, rootFiles []walker.Entry, commit *git.CommitInfo, instContent string, treeRoot *tree.Node, cfg *config.Config, picks []PickInfo) error {
	f, err := os.Create(path)
	if err != nil {
		return fmt.Errorf("cannot create root index: %w", err)
	}
	defer f.Close()

	var sb strings.Builder
	fmt.Fprintf(&sb, "# %s [ROOT]\n", absRoot)
	fmt.Fprintf(&sb, "> Generated: %s\n", time.Now().Format(time.RFC3339))
	sb.WriteString(git.FormatCommitLine(commit))
	if instContent != "" {
		sb.WriteString(instructions.Format(instContent))
	} else {
		sb.WriteString("\n")
	}
	sb.WriteString(FormatFlags(cfg))

	sb.WriteString("## Directory Tree\n\n")
	sb.WriteString(fenceChars + "\n")
	sb.WriteString(renderTreeWithLegend(treeRoot, rootName))
	sb.WriteString(fenceChars + "\n\n")

	sb.WriteString("## Vertical Slices Index\n\n")
	sb.WriteString("These files contain the full contents of each top-level directory.\n")
	sb.WriteString("Read the relevant slice file for detailed source code.\n\n")
	sb.WriteString("| Slice File | Directory | Files | Total Bytes |\n")
	sb.WriteString("|------------|-----------|-------|-------------|\n")

	rootSliceFile := "root_" + repoName + ".md"
	rootFileCount := len(rootFiles)
	rootTotalBytes := int64(0)
	for _, e := range rootFiles {
		rootTotalBytes += e.Size
	}
	fmt.Fprintf(&sb, "| %s | (root) | %d | %d |\n", rootSliceFile, rootFileCount, rootTotalBytes)

	for _, dirName := range sortedKeys(topLevelDirs) {
		entries := topLevelDirs[dirName]
		fileCount := len(entries)
		totalBytes := int64(0)
		for _, e := range entries {
			totalBytes += e.Size
		}
		sliceFile := filepath.Join(dirName, dirName+"_"+repoName+".md")
		fmt.Fprintf(&sb, "| %s | %s/ | %d | %d |\n", sliceFile, dirName, fileCount, totalBytes)
	}
	sb.WriteString("\n")

	if len(picks) > 0 {
		sb.WriteString("## Picks\n\n")
		sb.WriteString("These files were explicitly selected via --pick flags.\n\n")
		for _, p := range picks {
			fmt.Fprintf(&sb, "- `%s` (%s)\n", p.RelPath, p.SliceDir)
		}
		sb.WriteString("\n")
	}

	_, err = f.WriteString(sb.String())
	return err
}

func writeRootSliceFile(path, absRoot string, rootFiles []walker.Entry, omitTest bool, commit *git.CommitInfo, cfg *config.Config) error {
	f, err := os.Create(path)
	if err != nil {
		return fmt.Errorf("cannot create root slice file: %w", err)
	}
	defer f.Close()

	var sb strings.Builder
	fmt.Fprintf(&sb, "# %s [ROOT]\n", absRoot)
	fmt.Fprintf(&sb, "> Generated: %s\n", time.Now().Format(time.RFC3339))
	sb.WriteString(git.FormatCommitLine(commit))
	sb.WriteString(FormatFlags(cfg))

	rootName := filepath.Base(absRoot)
	fmt.Fprintf(&sb, "Tree: %s /\n\n", rootName)

	if len(rootFiles) == 0 {
		sb.WriteString("(none)\n")
		_, err = f.WriteString(sb.String())
		return err
	}

	sort.Slice(rootFiles, func(i, j int) bool { return rootFiles[i].RelPath < rootFiles[j].RelPath })
	for _, entry := range rootFiles {
		fullPath := filepath.Join(absRoot, entry.RelPath)
		content, err := readFileForOutput(fullPath, omitTest)
		if err != nil {
			fmt.Fprintf(os.Stderr, "warning: cannot read %s: %v\n", entry.RelPath, err)
			continue
		}
		language := extToLang(filepath.Ext(entry.RelPath))
		lines := countLines(content)
		fmt.Fprintf(&sb, "--- FILE: ./%s (%d bytes, %d lines) ---\n", entry.RelPath, len(content), lines)
		fmt.Fprintf(&sb, "%s", fenceChars)
		if language != "" {
			fmt.Fprintf(&sb, "%s", language)
		}
		fmt.Fprintf(&sb, "\n%s", content)
		if !strings.HasSuffix(content, "\n") {
			fmt.Fprintf(&sb, "\n")
		}
		fmt.Fprintf(&sb, "%s\n--- END FILE ---\n\n", fenceChars)
	}

	_, err = f.WriteString(sb.String())
	return err
}

func writeDirSlice(path, absDir, absRoot, dirName string, fileEntries []walker.Entry, rootFileName string, omitTest bool, commit *git.CommitInfo, cfg *config.Config) error {
	f, err := os.Create(path)
	if err != nil {
		return fmt.Errorf("cannot create slice file: %w", err)
	}
	defer f.Close()

	var sb strings.Builder
	fmt.Fprintf(&sb, "# %s [SLICE]\n", absDir)
	fmt.Fprintf(&sb, "> Generated: %s\n", time.Now().Format(time.RFC3339))
	sb.WriteString(git.FormatCommitLine(commit))
	sb.WriteString(FormatFlags(cfg))
	fmt.Fprintf(&sb, "Root document: %s\n\n", rootFileName)

	fmt.Fprintf(&sb, "Tree: %s / %s\n\n", filepath.Base(absRoot), dirName)

	if len(fileEntries) > 0 {
		sb.WriteString("## Contents\n\n")
		sort.Slice(fileEntries, func(i, j int) bool { return fileEntries[i].RelPath < fileEntries[j].RelPath })
		for _, entry := range fileEntries {
			fullPath := filepath.Join(absRoot, entry.RelPath)
			content, err := readFileForOutput(fullPath, omitTest)
			if err != nil {
				fmt.Fprintf(os.Stderr, "warning: cannot read %s: %v\n", entry.RelPath, err)
				continue
			}
			language := extToLang(filepath.Ext(entry.RelPath))
			lines := countLines(content)
			fmt.Fprintf(&sb, "--- FILE: ./%s (%d bytes, %d lines) ---\n", entry.RelPath, len(content), lines)
			fmt.Fprintf(&sb, "%s", fenceChars)
			if language != "" {
				fmt.Fprintf(&sb, "%s", language)
			}
			fmt.Fprintf(&sb, "\n%s", content)
			if !strings.HasSuffix(content, "\n") {
				fmt.Fprintf(&sb, "\n")
			}
			fmt.Fprintf(&sb, "%s\n--- END FILE ---\n\n", fenceChars)
		}
	}

	_, err = f.WriteString(sb.String())
	return err
}

func sortedKeys(m map[string][]walker.Entry) []string {
	keys := make([]string, 0, len(m))
	for k := range m {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	return keys
}

func renderTreeWithLegend(root *tree.Node, rootName string) string {
	var sb strings.Builder
	sb.WriteString(rootName)
	sb.WriteString("/\n")
	renderNodeLLM(&sb, root, "")

	if hasAnyAnnotation(root) {
		sb.WriteString("Legend: ~ ignored | ~ tests omitted | ~ omitted | ~ exceeded --max-depth\n")
	}

	return sb.String()
}

func renderNodeLLM(sb *strings.Builder, node *tree.Node, prefix string) {
	keys := make([]string, 0, len(node.Children))
	for k := range node.Children {
		keys = append(keys, k)
	}
	sort.Strings(keys)

	for i, key := range keys {
		child := node.Children[key]
		isLast := i == len(keys)-1

		sb.WriteString(prefix)
		sb.WriteString("+- ")

		sb.WriteString(child.Name)
		if child.IsDir {
			sb.WriteString("/")
		}
		if child.Annotation != "" {
			sb.WriteString(" ~ ")
			sb.WriteString(child.Annotation)
		}
		if child.Ref != "" {
			sb.WriteString(" -> see ")
			sb.WriteString(child.Ref)
		}
		sb.WriteString("\n")

		if child.IsDir && child.Children != nil {
			newPrefix := prefix
			if isLast {
				newPrefix += "    "
			} else {
				newPrefix += "|   "
			}
			renderNodeLLM(sb, child, newPrefix)
		}
	}
}

func hasAnyAnnotation(node *tree.Node) bool {
	if node.Annotation != "" {
		return true
	}
	for _, child := range node.Children {
		if hasAnyAnnotation(child) {
			return true
		}
	}
	return false
}

func readFile(path string) (string, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return "", err
	}

	for i, b := range data {
		if b == 0 {
			return "", fmt.Errorf("contains null byte at position %d", i)
		}
	}

	if !utf8.Valid(data) {
		return "", fmt.Errorf("invalid UTF-8 encoding")
	}

	content := string(data)
	content = strings.ReplaceAll(content, "\r\n", "\n")
	content = strings.ReplaceAll(content, "\r", "\n")

	return content, nil
}

func readFileForOutput(fullPath string, omitTest bool) (string, error) {
	content, err := readFile(fullPath)
	if err != nil {
		return "", err
	}
	if omitTest && strings.EqualFold(filepath.Ext(fullPath), ".rs") {
		content = stripTestBlocks(content)
	}
	return content, nil
}

func stripTestBlocks(content string) string {
	lines := strings.Split(content, "\n")
	var result []string

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

		if isCfgTestAttr(trimmed) {
			braceDepth := 0
			foundBrace := false

			for i < len(lines) {
				line := lines[i]
				braceDepth += strings.Count(line, "{") - strings.Count(line, "}")
				if strings.Contains(line, "{") {
					foundBrace = true
				}
				i++

				if foundBrace && braceDepth <= 0 {
					break
				}
			}
			continue
		}

		result = append(result, line)
		i++
	}

	if len(result) == 0 {
		return ""
	}
	return strings.Join(result, "\n")
}

func isCfgTestAttr(line string) bool {
	return strings.Contains(line, "#[cfg(test)]") ||
		strings.Contains(line, "#[cfg(all(test") ||
		strings.Contains(line, "#[cfg(any(test")
}

func extToLang(ext string) string {
	mapping := map[string]string{
		".go":   "go",
		".md":   "markdown",
		".yml":  "yaml",
		".yaml": "yaml",
		".js":   "javascript",
		".ts":   "typescript",
		".css":  "css",
		".html": "html",
		".json": "json",
		".xml":  "xml",
		".sql":  "sql",
		".sh":   "bash",
		".py":   "python",
		".rs":   "rust",
		".c":    "c",
		".h":    "c",
		".cpp":  "cpp",
		".hpp":  "cpp",
		".txt":  "text",
		".toml": "toml",
	}

	if lang, ok := mapping[strings.ToLower(ext)]; ok {
		return lang
	}
	return ""
}

func countLines(content string) int {
	if content == "" {
		return 0
	}
	count := strings.Count(content, "\n")
	if len(content) > 0 && content[len(content)-1] != '\n' {
		count++
	}
	return count
}