Files

483 lines
12 KiB
Go

package renderer
import (
"fmt"
"os"
"path/filepath"
"sort"
"strings"
"unicode/utf8"
"dirmd/internal/tree"
"dirmd/internal/walker"
)
const fenceChars = "~~~"
// Write writes a single markdown file with tree and file contents.
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\n", h1Path)
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 := readFile(fullPath)
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
}
// WriteSlices generates a root markdown file plus one markdown per top-level directory.
// Returns the list of generated file paths.
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
// Write root file as index_<reponame>.md
indexBasename := "index_" + repoName + ".md"
rootPath := filepath.Join(outputDir, indexBasename)
generated = append(generated, rootPath)
if err := writeRootSlice(rootPath, absRoot, rootName, repoName, rootFiles, topLevelDirs); err != nil {
return nil, fmt.Errorf("failed to write root slice: %w", err)
}
// Write each top-level directory slice into a mirroring subdirectory
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]
// Strip dirName prefix for tree building only
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); 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) 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\n", absRoot)
// Build combined tree with references
treeLines := buildTreeWithRefs(rootName, repoName, topLevelDirs, rootFiles)
sb.WriteString("## Directory Tree\n\n")
sb.WriteString(fenceChars + "\n")
sb.WriteString(treeLines)
sb.WriteString(fenceChars + "\n\n")
// Vertical slices index for LLM navigation
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 := readFile(fullPath)
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 writeDirSlice(path, absDir, absRoot, dirName string, treeEntries, fileEntries []walker.Entry, rootFileName string) 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\n", absDir)
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 := readFile(fullPath)
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 buildTreeWithRefs(rootName, repoName string, topLevelDirs map[string][]walker.Entry, rootFiles []walker.Entry) string {
var sb strings.Builder
sb.WriteString(rootName)
sb.WriteString("/\n")
// Collect and sort directory names
var dirNames []string
for k := range topLevelDirs {
dirNames = append(dirNames, k)
}
sort.Strings(dirNames)
// Sort root files
sort.Slice(rootFiles, func(i, j int) bool {
return rootFiles[i].RelPath < rootFiles[j].RelPath
})
// Merge directories and files into a single sorted list for display
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 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
}