package llm

import "strings"

// StripMarkdownCodeBlocks removes markdown code block markers (```json ... ```) from LLM responses
// This handles cases where Claude returns JSON wrapped in markdown formatting
// If multiple code blocks exist, extracts the last one (assumes LLM's final answer)
func StripMarkdownCodeBlocks(response string) string {
	response = strings.TrimSpace(response)

	// Find all code blocks and extract the last valid JSON one
	// Pattern: ```json ... ``` or ``` ... ```
	blocks := extractCodeBlocks(response)
	if len(blocks) > 0 {
		// Return the last block (LLM's final answer)
		return blocks[len(blocks)-1]
	}

	// No code blocks found, return as-is
	return response
}

// extractCodeBlocks finds all markdown code blocks in the response
func extractCodeBlocks(text string) []string {
	var blocks []string
	lines := strings.Split(text, "\n")

	var currentBlock strings.Builder
	inBlock := false

	for _, line := range lines {
		trimmed := strings.TrimSpace(line)

		// Check if this line is a code block delimiter
		if trimmed == "```json" || trimmed == "```" {
			if inBlock {
				// Closing a block
				content := strings.TrimSpace(currentBlock.String())
				if content != "" {
					blocks = append(blocks, content)
				}
				currentBlock.Reset()
				inBlock = false
			} else {
				// Opening a block
				inBlock = true
			}
		} else if inBlock {
			// Inside a block, accumulate content
			if currentBlock.Len() > 0 {
				currentBlock.WriteString("\n")
			}
			currentBlock.WriteString(line)
		}
	}

	// Handle case where block wasn't closed
	if inBlock && currentBlock.Len() > 0 {
		content := strings.TrimSpace(currentBlock.String())
		if content != "" {
			blocks = append(blocks, content)
		}
	}

	return blocks
}
