ctx/output.go

package main

import (
	"bufio"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"slices"
	"strings"
)

func writeOutput(root string, files []string, outputPath string) (count int64, err error) {
	f, err := os.Create(outputPath)
	if err != nil {
		return 0, err
	}
	defer func() {
		if cerr := f.Close(); cerr != nil && err == nil {
			err = cerr
		}
	}()

	bw := bufio.NewWriterSize(f, 1024*1024)
	defer func() {
		if ferr := bw.Flush(); ferr != nil && err == nil {
			err = ferr
		}
	}()

	tc := &TokenCounter{w: bw}

	fmt.Fprintf(tc, "Project Path: %s\n\n", filepath.Base(root))
	fmt.Fprintln(tc, "Source Tree:")
	fmt.Fprintln(tc, "")

	fmt.Fprintln(tc, "```txt")
	fmt.Fprintln(tc, filepath.Base(root))

	if err := writeTree(tc, files); err != nil {
		return 0, err
	}

	fmt.Fprintln(tc, "```")
	fmt.Fprintln(tc, "")

	for _, file := range files {
		if file == outputPath || filepath.Base(file) == outputPath {
			continue
		}

		fullPath := filepath.Join(root, file)
		content, err := os.ReadFile(fullPath)
		if err != nil {
			fmt.Fprintf(tc, "Error reading %s: %v\n", file, err)
			continue
		}

		ext := strings.TrimPrefix(filepath.Ext(file), ".")
		if ext == "" {
			ext = "txt"
		}

		fmt.Fprintf(tc, "`%s`:\n\n", file)
		fmt.Fprintf(tc, "```%s\n", ext)

		if _, err := tc.Write(content); err != nil {
			return 0, err
		}

		if len(content) > 0 && content[len(content)-1] != '\n' {
			if err := tc.WriteByte('\n'); err != nil {
				return 0, err
			}
		}
		fmt.Fprintln(tc, "```")
		fmt.Fprintln(tc, "")
	}

	return tc.Count, tc.Err
}

func writeTree(w io.Writer, files []string) error {
	root := make(map[string]any)
	for _, f := range files {
		parts := strings.Split(filepath.ToSlash(f), "/")
		current := root
		for _, part := range parts {
			if _, exists := current[part]; !exists {
				current[part] = make(map[string]any)
			}
			current = current[part].(map[string]any)
		}
	}

	return printNode(w, root, "")
}

func printNode(w io.Writer, node map[string]any, prefix string) error {
	keys := make([]string, 0, len(node))
	for k := range node {
		keys = append(keys, k)
	}
	slices.Sort(keys)

	for i, key := range keys {
		isLast := i == len(keys)-1
		connector := "├── "
		if isLast {
			connector = "└── "
		}

		if _, err := fmt.Fprintf(w, "%s%s%s\n", prefix, connector, key); err != nil {
			return err
		}

		children := node[key].(map[string]any)
		if len(children) > 0 {
			childPrefix := prefix + "│   "
			if isLast {
				childPrefix = prefix + "    "
			}
			if err := printNode(w, children, childPrefix); err != nil {
				return err
			}
		}
	}
	return nil
}