refactor(tui): restructure code highlighting with better separation of concerns

- Extract `parseRange` function to handle range parsing logic
- Split `renderWithStyle` into `renderHighlightedCode` and `renderPlainCode` functions
- Add `formatLineNumber` and `getLineNumberWidth` helper functions for line number formatting
- Extract `parseCodeBlockInfo` and `renderMarkdownSection` functions from main processing loop
- Improve error handling with fallback to plain rendering when syntax highlighting fails
- Update tests to reflect new function names and signatures
- Simplify regex pattern construction in `processMarkdownWithHighlighting`

Author: tharropoulos

internal/tui/highlight.go

@@ -33,26 +33,17 @@ func parseHighlightSyntax(syntax string) []LineRange {
 	}
 
 	syntax = strings.Trim(syntax, "{}")
-
-	var ranges []LineRange
 	parts := strings.Split(syntax, ",")
+	var ranges []LineRange
 
 	for _, part := range parts {
 		part = strings.TrimSpace(part)
 		if strings.Contains(part, "-") {
-			// Range like "1-3"
-			rangeParts := strings.Split(part, "-")
-			if len(rangeParts) == 2 {
-				start, err1 := strconv.Atoi(strings.TrimSpace(rangeParts[0]))
-				end, err2 := strconv.Atoi(strings.TrimSpace(rangeParts[1]))
-				if err1 == nil && err2 == nil && start <= end {
-					ranges = append(ranges, LineRange{Start: start, End: end})
-				}
+			if r := parseRange(part); r != nil {
+				ranges = append(ranges, *r)
 			}
 		} else {
-			// Single line like "5"
-			line, err := strconv.Atoi(part)
-			if err == nil {
+			if line, err := strconv.Atoi(part); err == nil {
 				ranges = append(ranges, LineRange{Start: line, End: line})
 			}
 		}
@@ -61,7 +52,26 @@ func parseHighlightSyntax(syntax string) []LineRange {
 	return ranges
 }
 
+func parseRange(part string) *LineRange {
+	rangeParts := strings.Split(part, "-")
+	if len(rangeParts) != 2 {
+		return nil
+	}
+
+	start, err1 := strconv.Atoi(strings.TrimSpace(rangeParts[0]))
+	end, err2 := strconv.Atoi(strings.TrimSpace(rangeParts[1]))
+
+	if err1 == nil && err2 == nil && start <= end {
+		return &LineRange{Start: start, End: end}
+	}
+	return nil
+}
+
 func shouldHighlightLine(lineNum int, ranges []LineRange) bool {
+	if len(ranges) == 0 {
+		return true // highlight all lines if no ranges specified
+	}
+
 	for _, r := range ranges {
 		if lineNum >= r.Start && lineNum <= r.End {
 			return true
@@ -70,102 +80,185 @@ func shouldHighlightLine(lineNum int, ranges []LineRange) bool {
 	return false
 }
 
-func renderCustomCodeBlock(content string, info CodeHighlightInfo, themeName string) string {
-	lexer := lexers.Get(info.Language)
-	if lexer == nil {
-		lexer = lexers.Fallback
+func formatLineNumber(lineNum, width int) string {
+	style := lipgloss.NewStyle().
+		Foreground(lipgloss.Color("240")).
+		PaddingRight(1)
+
+	lineNumStr := strconv.Itoa(lineNum)
+	paddedLineNum := fmt.Sprintf("%*s", width-1, lineNumStr)
+	return style.Render(paddedLineNum)
+}
+
+func getLineNumberWidth(startLine, totalLines int) int {
+	if totalLines == 0 {
+		return 0
 	}
-	lexer = chroma.Coalesce(lexer)
+	maxLineNum := startLine + totalLines - 1
+	return len(strconv.Itoa(maxLineNum)) + 2
+}
 
-	style := config.GetChromaStyle(themeName)
-	return renderWithStyle(content, info, lexer, style)
+func renderPlainCode(lines []string, info CodeHighlightInfo) string {
+	var result strings.Builder
+	lineNumberWidth := 0
+
+	if info.ShowLineNumbers {
+		lineNumberWidth = getLineNumberWidth(info.StartLine, len(lines))
+	}
+
+	for i, line := range lines {
+		displayLineNum := info.StartLine + i
+
+		if info.ShowLineNumbers {
+			result.WriteString(formatLineNumber(displayLineNum, lineNumberWidth))
+		}
+
+		result.WriteString(line)
+
+		if i < len(lines)-1 {
+			result.WriteString("\n")
+		}
+	}
+
+	return result.String()
 }
 
-func renderWithStyle(content string, info CodeHighlightInfo, lexer chroma.Lexer, style *chroma.Style) string {
+func renderHighlightedCode(content string, lines []string, info CodeHighlightInfo, lexer chroma.Lexer, style *chroma.Style) string {
 	formatter := formatters.Get("terminal256")
 	if formatter == nil {
-		formatter = formatters.Fallback
+		return renderPlainCode(lines, info)
 	}
 
-	lines := strings.Split(content, "\n")
+	iterator, err := lexer.Tokenise(nil, content)
+	if err != nil {
+		return renderPlainCode(lines, info)
+	}
+
+	var formattedBuf strings.Builder
+	if err := formatter.Format(&formattedBuf, style, iterator); err != nil {
+		return renderPlainCode(lines, info)
+	}
 
-	highlightStyle := lipgloss.NewStyle()
+	formattedLines := strings.Split(formattedBuf.String(), "\n")
 
+	// Ensure we have the same number of lines
+	for len(formattedLines) < len(lines) {
+		formattedLines = append(formattedLines, "")
+	}
+	if len(formattedLines) > len(lines) {
+		formattedLines = formattedLines[:len(lines)]
+	}
+
+	var result strings.Builder
 	lineNumberWidth := 0
+
 	if info.ShowLineNumbers {
-		maxLineNum := info.StartLine + len(lines) - 1
-		lineNumberWidth = len(strconv.Itoa(maxLineNum)) + 2
+		lineNumberWidth = getLineNumberWidth(info.StartLine, len(lines))
 	}
 
-	lineNumberStyle := lipgloss.NewStyle().
-		Foreground(lipgloss.Color("240")).
-		PaddingRight(1)
+	for i, line := range lines {
+		displayLineNum := info.StartLine + i
+		relativeLineNum := i + 1
 
-	var result strings.Builder
-	for lineNum, line := range lines {
-		displayLineNum := info.StartLine + lineNum
-		relativeLineNum := lineNum + 1
-
-		// Add line number if requested
 		if info.ShowLineNumbers {
-			lineNumStr := strconv.Itoa(displayLineNum)
-			paddedLineNum := fmt.Sprintf("%*s", lineNumberWidth-1, lineNumStr)
-			result.WriteString(lineNumberStyle.Render(paddedLineNum))
+			result.WriteString(formatLineNumber(displayLineNum, lineNumberWidth))
 		}
 
-		shouldHighlight := len(info.Ranges) == 0 || shouldHighlightLine(relativeLineNum, info.Ranges)
-
-		if !shouldHighlight {
-			result.WriteString(line)
-			if lineNum < len(lines)-1 {
-				result.WriteString("\n")
+		if shouldHighlightLine(relativeLineNum, info.Ranges) {
+			formattedLine := ""
+			if i < len(formattedLines) {
+				formattedLine = strings.TrimRight(formattedLines[i], " \t\n\r")
 			}
-			continue
-		}
-
-		lineIterator, err := lexer.Tokenise(nil, line)
-		if err != nil {
-			result.WriteString(highlightStyle.Render(line))
-			if lineNum < len(lines)-1 {
-				result.WriteString("\n")
+			if formattedLine == "" {
+				formattedLine = line
 			}
-			continue
+			result.WriteString(formattedLine)
+		} else {
+			result.WriteString(line)
 		}
 
-		var lineBuf strings.Builder
-		err = formatter.Format(&lineBuf, style, lineIterator)
-		if err != nil {
-			result.WriteString(highlightStyle.Render(line))
-			if lineNum < len(lines)-1 {
-				result.WriteString("\n")
-			}
-			continue
+		if i < len(lines)-1 {
+			result.WriteString("\n")
 		}
+	}
 
-		syntaxHighlighted := lineBuf.String()
-		syntaxHighlighted = strings.TrimRight(syntaxHighlighted, " \t\n\r")
-		result.WriteString(highlightStyle.Render(syntaxHighlighted))
+	return result.String()
+}
 
-		if lineNum < len(lines)-1 {
-			result.WriteString("\n")
+func renderCustomCodeBlock(content string, info CodeHighlightInfo, themeName string) string {
+	lines := strings.Split(content, "\n")
+
+	var renderedContent string
+
+	if info.Language != "" {
+		lexer := lexers.Get(info.Language)
+		if lexer == nil {
+			lexer = lexers.Fallback
 		}
+		lexer = chroma.Coalesce(lexer)
+		style := config.GetChromaStyle(themeName)
+
+		renderedContent = renderHighlightedCode(content, lines, info, lexer, style)
+	} else {
+		renderedContent = renderPlainCode(lines, info)
 	}
 
+	// Apply consistent styling
 	codeStyle := lipgloss.NewStyle().
 		Padding(1, 2).
 		Width(78).
 		MarginTop(1).
 		MarginBottom(1)
 
-	return codeStyle.Render(result.String())
+	return codeStyle.Render(renderedContent)
+}
+
+func parseCodeBlockInfo(match []string) CodeHighlightInfo {
+	if len(match) < 7 {
+		return CodeHighlightInfo{}
+	}
+
+	language := match[1]
+	highlightSyntax := ""
+	if len(match) > 2 && match[2] != "" {
+		highlightSyntax = match[2]
+	}
+	numberedFlag := strings.TrimSpace(match[3])
+	startLineStr := match[5]
+
+	startLine := 1
+	if startLineStr != "" {
+		if parsed, err := strconv.Atoi(startLineStr); err == nil && parsed > 0 {
+			startLine = parsed
+		}
+	}
+
+	return CodeHighlightInfo{
+		Language:        language,
+		Ranges:          parseHighlightSyntax(highlightSyntax),
+		ShowLineNumbers: strings.Contains(numberedFlag, "--numbered"),
+		StartLine:       startLine,
+	}
+}
+
+func renderMarkdownSection(text, themeName string) string {
+	if strings.TrimSpace(text) == "" {
+		return ""
+	}
+
+	rendered, err := glamour.Render(text, themeName)
+	if err != nil {
+		return text
+	}
+	return rendered
 }
 
 func processMarkdownWithHighlighting(markdown string, themeName string) (string, error) {
-	re := regexp.MustCompile("(?s)```([a-zA-Z0-9_+-]*)({[^}]*})?(\\s+--numbered)?(\\s+--start-at-line\\s+(\\d+))?\\s*\n(.*?)\n```")
+	// Regex to match code blocks with optional highlighting syntax
+	re := regexp.MustCompile(`(?s)` + "`" + `{3}([a-zA-Z0-9_+-]*)({[^}]*})?(\s+--numbered)?(\s+--start-at-line\s+(\d+))?\s*\n(.*?)\n` + "`" + `{3}`)
 
 	matches := re.FindAllStringSubmatch(markdown, -1)
 	if len(matches) == 0 {
-		// No code blocks found, render the entire markdown with Glamour
 		return glamour.Render(markdown, themeName)
 	}
 
@@ -174,80 +267,35 @@ func processMarkdownWithHighlighting(markdown string, themeName string) (string,
 	indices := re.FindAllStringIndex(markdown, -1)
 
 	for i, match := range matches {
-		if len(match) < 7 {
-			continue
-		}
-
 		matchStart := indices[i][0]
 		matchEnd := indices[i][1]
 
-		// Add the text before this code block
+		// Add text before this code block
 		if matchStart > lastIndex {
 			beforeText := markdown[lastIndex:matchStart]
-			if strings.TrimSpace(beforeText) != "" {
-				// Render regular markdown content with Glamour
-				rendered, err := glamour.Render(beforeText, themeName)
-				if err != nil {
-					parts = append(parts, beforeText)
-				} else {
-					parts = append(parts, rendered)
-				}
-			}
+			parts = append(parts, renderMarkdownSection(beforeText, themeName))
 		}
 
-		// Extract code block info
-		language := match[1] // e.g., "typescript"
-		highlightSyntax := ""
-		if len(match) > 2 && match[2] != "" {
-			highlightSyntax = match[2] // e.g., "{1-2}"
-		}
-		numberedFlag := strings.TrimSpace(match[3]) // " --numbered"
-		startLineStr := match[5]                    // The line number from --start-at-line
-		content := match[6]                         // The actual code content
-
-		startLine := 1
-		if startLineStr != "" {
-			if parsed, err := strconv.Atoi(startLineStr); err == nil && parsed > 0 {
-				startLine = parsed
-			}
-		}
+		// Process the code block
+		info := parseCodeBlockInfo(match)
+		content := match[6]
 
-		info := CodeHighlightInfo{
-			Language:        language,
-			Ranges:          parseHighlightSyntax(highlightSyntax),
-			ShowLineNumbers: strings.Contains(numberedFlag, "--numbered"),
-			StartLine:       startLine,
-		}
-
-		if language != "" || info.ShowLineNumbers {
+		if info.Language != "" || info.ShowLineNumbers {
 			customRendered := renderCustomCodeBlock(content, info, themeName)
 			parts = append(parts, customRendered)
 		} else {
-			// No language specified and no line numbers, use Glamour directly for the entire code block
+			// Use Glamour for plain code blocks
 			codeBlock := match[0]
-			rendered, err := glamour.Render(codeBlock, themeName)
-			if err != nil {
-				parts = append(parts, codeBlock)
-			} else {
-				parts = append(parts, rendered)
-			}
+			parts = append(parts, renderMarkdownSection(codeBlock, themeName))
 		}
 
 		lastIndex = matchEnd
 	}
 
-	// Add any remaining text after the last code block
+	// Add remaining text after the last code block
 	if lastIndex < len(markdown) {
 		remainingText := markdown[lastIndex:]
-		if strings.TrimSpace(remainingText) != "" {
-			// Render remaining markdown content with Glamour
-			rendered, err := glamour.Render(remainingText, themeName)
-			if err != nil {
-				parts = append(parts, remainingText)
-			} else {
-				parts = append(parts, rendered)
-			}
-		}
+		parts = append(parts, renderMarkdownSection(remainingText, themeName))
 	}
 
 	return strings.Join(parts, ""), nil