main.go

package main

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

var styles = map[string]Style{
	".bash": {LineComment: "#"},
	".css":  {BlockDo: "/*", BlockDone: "*/"},
	".go":   {LineComment: "//"},
	".haml": {LineComment: "-#"},
	".html": {BlockDo: "<!--", BlockDone: "-->"},
	".js":   {LineComment: "//", BlockDo: "/*", BlockDone: "*/"},
	".lua":  {LineComment: "--"},
	".rb":   {LineComment: "#"},
	".scss": {LineComment: "//", BlockDo: "/*", BlockDone: "*/"},
	".sh":   {LineComment: "#"},
	".sql":  {LineComment: "--"},
	".ts":   {LineComment: "//", BlockDo: "/*", BlockDone: "*/"},
}

type ProcessState struct {
	filesInProcess map[string]bool
}

type Style struct {
	LineComment string
	BlockDo     string
	BlockDone   string
}

func main() {
	state := &ProcessState{
		filesInProcess: make(map[string]bool),
	}
	if err := processMD(os.Stdin, os.Stdout, state); err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
		os.Exit(1)
	}
}

// processMD reads Markdown from input and writes to output, processing embed blocks
func processMD(input io.Reader, output io.Writer, state *ProcessState) error {
	scanner := bufio.NewScanner(input)
	inEmbedBlock := false // Flag to track if we're inside an embed block
	var lines []string    // Collects lines within an embed block

	for scanner.Scan() {
		line := scanner.Text()

		if !inEmbedBlock {
			if line == "```embed" {
				// Start of embed block
				inEmbedBlock = true
				lines = []string{}
			} else {
				// Write line directly to output
				fmt.Fprintln(output, line)
			}
		} else {
			if line == "```" {
				// End of an embed block
				if err := processEmbed(lines, output, state); err != nil {
					return err
				}
				inEmbedBlock = false
			} else {
				// Collect lines in embed block
				lines = append(lines, line)
			}
		}
	}

	if inEmbedBlock {
		return fmt.Errorf("unterminated ```embed code block")
	}

	if err := scanner.Err(); err != nil {
		return err
	}

	return nil
}

// processEmbed processes lines collected within an embed block
func processEmbed(lines []string, output io.Writer, state *ProcessState) error {
	for i, line := range lines {
		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}

		parts := strings.Fields(line)
		if len(parts) == 0 {
			continue
		}

		filename := parts[0] // Required filename
		blockName := ""      // Optional block name

		if len(parts) == 2 {
			blockName = parts[1]
		} else if len(parts) > 2 {
			return fmt.Errorf("invalid format in embed code block: %s", line)
		}

		content, err := os.ReadFile(filename)
		if err != nil {
			return fmt.Errorf("failed to read file %s: %v", filename, err)
		}
		fileContent := string(content)

		if err := processFile(filename, blockName, fileContent, output, state); err != nil {
			return err
		}

		// Add newline between multiple code blocks
		if i < len(lines)-1 {
			fmt.Fprintln(output)
		}
	}
	return nil
}

// processFile processes an individual file, handling circular embeddings
func processFile(filename, blockName, fileContent string, output io.Writer, state *ProcessState) error {
	if state.filesInProcess[filename] {
		return fmt.Errorf("circular embedding detected for file %s", filename)
	}

	// Mark the file as being processed
	state.filesInProcess[filename] = true
	defer delete(state.filesInProcess, filename)

	ext := filepath.Ext(filename)
	if ext == ".md" {
		// Process Markdown files recursively
		reader := strings.NewReader(fileContent)
		if err := processMD(reader, output, state); err != nil {
			return fmt.Errorf("processing markdown file %s failed: %v", filename, err)
		}
	} else {
		// Process other files as code blocks
		if err := processCodeFile(filename, blockName, fileContent, output); err != nil {
			return err
		}
	}
	return nil
}

// processCodeFile processes non-Markdown files and embeds their content in code fences
func processCodeFile(filename, blockName, fileContent string, output io.Writer) error {
	ext := filepath.Ext(filename)
	lang := strings.TrimPrefix(ext, ".")

	// Get comment style based on file extension
	style, ok := styles[ext]
	if !ok {
		return fmt.Errorf("unsupported file type: %s", ext)
	}

	// Prepare filename comment
	var fileName string
	if style.LineComment != "" {
		fileName = style.LineComment + " " + filename
	} else if style.BlockDo != "" && style.BlockDone != "" {
		fileName = fmt.Sprintf("%s %s %s", style.BlockDo, filename, style.BlockDone)
	}

	// If a block name is specified, extract block between marks
	if blockName != "" {
		doMark, doneMark := getBlockMarkers(style, blockName)

		extractedContent, err := extractBlock(fileContent, doMark, doneMark)
		if err != nil {
			return fmt.Errorf("%v in file %s", err, filename)
		}

		fileContent = extractedContent
	}

	// Clean up content
	fileContent = strings.Trim(fileContent, "\n")
	fileContent = dedent(fileContent)

	// Write code block to output
	fmt.Fprintf(output, "```%s\n", lang)
	fmt.Fprintf(output, "%s\n", fileName)
	fmt.Fprintf(output, "%s", fileContent)
	fmt.Fprintf(output, "\n```\n")

	return nil
}

// getBlockMarkers generates the start and end markers for a block in a file
func getBlockMarkers(style Style, blockName string) (string, string) {
	blockName = strings.TrimSpace(blockName)
	var doMark, doneMark string

	if style.LineComment != "" {
		// Line comment marks
		doMark = strings.TrimSpace(fmt.Sprintf("%s emdo %s", style.LineComment, blockName))
		doneMark = strings.TrimSpace(fmt.Sprintf("%s emdone %s", style.LineComment, blockName))
	} else if style.BlockDo != "" && style.BlockDone != "" {
		// Block comment marks
		beginContent := strings.TrimSpace(fmt.Sprintf("emdo %s", blockName))
		endContent := strings.TrimSpace(fmt.Sprintf("emdone %s", blockName))

		doMark = fmt.Sprintf("%s %s %s", style.BlockDo, beginContent, style.BlockDone)
		doneMark = fmt.Sprintf("%s %s %s", style.BlockDo, endContent, style.BlockDone)
	} else {
		return "", ""
	}

	return doMark, doneMark
}

// extractBlock extracts the content between doMark and doneMark, ignoring leading and trailing whitespace
func extractBlock(fileContent, doMark, doneMark string) (string, error) {
	lines := strings.Split(fileContent, "\n")
	var inBlock bool
	var blockLines []string

	for _, line := range lines {
		trimmedLine := strings.TrimSpace(line)
		if !inBlock {
			if trimmedLine == doMark {
				inBlock = true
			}
		} else {
			if trimmedLine == doneMark {
				inBlock = false
				break
			}
			blockLines = append(blockLines, line)
		}
	}

	if inBlock {
		return "", fmt.Errorf("done mark '%s' not found", doneMark)
	}

	if len(blockLines) == 0 {
		return "", fmt.Errorf("no content found between do mark '%s' and done mark '%s'", doMark, doneMark)
	}

	return strings.Join(blockLines, "\n"), nil
}

// dedent removes common leading whitespace from all lines
func dedent(s string) string {
	lines := strings.Split(s, "\n")
	minIndent := -1

	// Find minimum indentation level
	for _, line := range lines {
		trimmed := strings.TrimLeft(line, " \t")
		if trimmed == "" {
			continue // Skip empty or whitespace-only lines
		}
		indent := len(line) - len(trimmed)
		if minIndent == -1 || indent < minIndent {
			minIndent = indent
		}
	}

	// Remove minimum indentation from each line
	if minIndent > 0 {
		for i, line := range lines {
			if len(line) >= minIndent {
				lines[i] = line[minIndent:]
			}
		}
	}

	return strings.Join(lines, "\n")
}