package rag

import "strings"

func ChunkMarkdown(content string, maxChunkSize int) []string {
	if maxChunkSize <= 0 {
		maxChunkSize = 500
	}

	sections := splitByHeaders(content)
	var chunks []string

	for _, section := range sections {
		if len(section) <= maxChunkSize {
			chunks = append(chunks, section)
			continue
		}
		paragraphs := splitByParagraphs(section)
		current := ""
		for _, p := range paragraphs {
			if len(current)+len(p)+2 > maxChunkSize && current != "" {
				chunks = append(chunks, strings.TrimSpace(current))
				current = ""
			}
			if current != "" {
				current += "\n\n"
			}
			current += p
		}
		if current != "" {
			chunks = append(chunks, strings.TrimSpace(current))
		}
	}
	return chunks
}

func splitByHeaders(content string) []string {
	lines := strings.Split(content, "\n")
	var sections []string
	current := ""

	for _, line := range lines {
		if strings.HasPrefix(line, "## ") && current != "" {
			sections = append(sections, strings.TrimSpace(current))
			current = ""
		}
		if current != "" {
			current += "\n"
		}
		current += line
	}
	if current != "" {
		sections = append(sections, strings.TrimSpace(current))
	}
	return sections
}

func splitByParagraphs(content string) []string {
	paragraphs := strings.Split(content, "\n\n")
	var result []string
	for _, p := range paragraphs {
		p = strings.TrimSpace(p)
		if p != "" {
			result = append(result, p)
		}
	}
	return result
}