dirmd/internal/renderer/renderer.go

package renderer

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

	"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")
		}
	}

	if !singleFile {
		root := tree.Build(result.Entries)
		rootName := filepath.Base(result.AbsRoot)
		sb.WriteString("## Directory Tree\n\n")
		sb.WriteString(fenceChars + "\n")
		sb.WriteString(renderTreeForLLM(root, 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, entry.Size, 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)
		}
	}

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

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

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

	for _, dirName := range dirNames {
		sliceName := strings.TrimSuffix(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]

		treeEntries := make([]walker.Entry, len(dirEntries))
		for i, e := range dirEntries {
			sepIdx := strings.Index(e.RelPath, string(filepath.Separator))
			if sepIdx >= 0 {
				treeEntries[i] = walker.Entry{RelPath: e.RelPath[sepIdx+1:], Size: e.Size}
			} else {
				treeEntries[i] = e
			}
		}

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

	return generated, nil
}

func writeRootSlice(path, absRoot, rootName, repoName string, rootFiles []walker.Entry, topLevelDirs map[string][]walker.Entry, omitTest bool, commit *git.CommitInfo, instContent string) error {
	f, err := os.Create(path)
	if err != nil {
		return fmt.Errorf("cannot create root 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))
	if instContent != "" {
		sb.WriteString(instructions.Format(instContent))
	} else {
		sb.WriteString("\n")
	}

	treeLines := buildTreeWithRefs(rootName, repoName, topLevelDirs, rootFiles)
	sb.WriteString("## Directory Tree\n\n")
	sb.WriteString(fenceChars + "\n")
	sb.WriteString(treeLines)
	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")

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

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

	if len(rootFiles) > 0 {
		sb.WriteString("## Root Contents\n\n")
		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, entry.Size, 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, treeEntries, fileEntries []walker.Entry, rootFileName string, omitTest bool, commit *git.CommitInfo) 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("\n")
	fmt.Fprintf(&sb, "Root document: %s\n\n", rootFileName)

	sb.WriteString("## Directory Tree\n\n")
	sb.WriteString(fenceChars + "\n")
	sliceTree := buildTreeFromEntries(dirName, treeEntries)
	sb.WriteString(sliceTree)
	sb.WriteString(fenceChars + "\n\n")

	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, entry.Size, 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 buildTreeWithRefs(rootName, repoName string, topLevelDirs map[string][]walker.Entry, rootFiles []walker.Entry) string {
	var sb strings.Builder
	sb.WriteString(rootName)
	sb.WriteString("/\n")

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

	sort.Slice(rootFiles, func(i, j int) bool {
		return rootFiles[i].RelPath < rootFiles[j].RelPath
	})

	type treeItem struct {
		name     string
		isDir    bool
		sliceRef string
	}

	items := make([]treeItem, 0, len(dirNames)+len(rootFiles))
	for _, d := range dirNames {
		sliceRef := filepath.Join(d, strings.TrimSuffix(d, "/")+"_"+repoName+".md")
		items = append(items, treeItem{
			name:     d,
			isDir:    true,
			sliceRef: sliceRef,
		})
	}
	for _, rf := range rootFiles {
		items = append(items, treeItem{
			name:  rf.RelPath,
			isDir: false,
		})
	}

	sort.Slice(items, func(i, j int) bool {
		return items[i].name < items[j].name
	})

	for _, item := range items {
		sb.WriteString("+- ")
		sb.WriteString(item.name)
		if item.isDir {
			sb.WriteString("/")
			sb.WriteString(" -> see ")
			sb.WriteString(item.sliceRef)
		}
		sb.WriteString("\n")
	}

	return sb.String()
}

func buildTreeFromEntries(rootName string, entries []walker.Entry) string {
	root := &tree.Node{IsDir: true, Children: make(map[string]*tree.Node)}

	for _, e := range entries {
		parts := strings.Split(e.RelPath, string(filepath.Separator))
		cur := root

		for i, part := range parts {
			isLeaf := i == len(parts)-1
			if _, ok := cur.Children[part]; !ok {
				node := &tree.Node{Name: part, IsDir: !isLeaf}
				if !isLeaf {
					node.Children = make(map[string]*tree.Node)
				}
				cur.Children[part] = node
			}
			cur = cur.Children[part]
		}
	}

	var sb strings.Builder
	sb.WriteString(rootName)
	sb.WriteString("/\n")
	renderNodeLLM(&sb, root, "")
	return sb.String()
}

func renderTreeForLLM(root *tree.Node, rootName string) string {
	var sb strings.Builder
	sb.WriteString(rootName)
	sb.WriteString("/\n")
	renderNodeLLM(&sb, root, "")
	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("/")
		}
		sb.WriteString("\n")

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

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]

				if !foundBrace {
					if strings.Contains(line, "{") {
						braceDepth += strings.Count(line, "{") - strings.Count(line, "}")
						foundBrace = true
					}
					i++
					continue
				}

				braceDepth += strings.Count(line, "{") - strings.Count(line, "}")
				i++

				if 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",
	}

	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
}