def _is_section(section): is_dict(section) && section[:type] == :section;

def _h_level_num(md):
  if (is_h1(md)): 1
  elif (is_h2(md)): 2
  elif (is_h3(md)): 3
  elif (is_h4(md)): 4
  elif (is_h5(md)): 5
  elif (is_h6(md)): 6
  else: 0
end

# Returns sections whose title contains the specified pattern.
# If depth is true, each section spans until the next header at the same or higher level.
def section(md_nodes, pattern, depth = false):
  if (depth):
    if (is_empty(md_nodes)):
      []
    else:
      do
        let n = len(md_nodes)
        | let match_indices = do foreach (i, range(n - 1)): if (is_h(md_nodes[i]) && contains(to_text(md_nodes[i]), pattern)): i; | compact();
        | var result = []
        | var i = 0
        | while (i < len(match_indices)):
            let start = match_indices[i]
            | let lvl = _h_level_num(md_nodes[start])
            | let boundaries = do foreach (j, range(n - 1)): if (is_h(md_nodes[j]) && _h_level_num(md_nodes[j]) <= lvl): j; | compact();
            | let boundaries_with_end = boundaries + n
            | let end_node = first(filter(boundaries_with_end, fn(b): b > start;))
            | result += [{type: :section, header: md_nodes[start], children: md_nodes[start + 1:end_node]}]
            | i += 1
            | result
          end
      end
  else:
    sections(md_nodes) | title_contains(pattern)
end

# Splits markdown nodes into sections based on headers.
def sections(md_nodes):
  if (is_empty(md_nodes)):
    []
  else:
    do
      let indices = do foreach (i, range(len(md_nodes) - 1)): if (is_h(md_nodes[i])): i; | compact();
      | let indices_with_end = indices + len(md_nodes)
      | let indices_len = len(indices)
      | var result = []
      | var i = 0
      | while (i < indices_len):
          let start_node = indices[i]
          | let end_node = indices_with_end[i + 1]
          | let children = md_nodes[start_node + 1:end_node]
          | result += [{type: :section, header: md_nodes[start_node], children: children}]
          | i = i + 1
          | result
        end
    end
end

# Filters sections based on a given predicate function.
def filter_sections(md_nodes, predicate):
  sections(md_nodes) | filter(fn(section): predicate(section);)
end

# Maps sections using a given mapper function.
def map_sections(md_nodes, mapper):
  sections(md_nodes) | map(fn(section): mapper(section[:header], section[:children]);)
end

# Returns an array of sections, each section is an array of markdown nodes between the specified header and the next header of the same level.
def split(md_nodes, level):
  if (is_empty(md_nodes)):
    []
  else:
    do
      let indices = do foreach (i, range(len(md_nodes) - 1)): if (is_h_level(md_nodes[i], level)): i; | compact();
      | let indices_with_end = indices + len(md_nodes)
      | let indices_len = len(indices)
      | var result = []
      | var i = 0
      | while (i < len(indices)):
          let start_node = indices[i]
          | let end_node = indices_with_end[i + 1]
          | let children = md_nodes[start_node + 1:end_node]
          | result += [{type: :section, header: md_nodes[start_node], children: children}]
          | i += 1
          | result
        end
    end
end

# Filters the given list of sections, returning only those whose title contains the specified text.
def title_contains(sections, text):
  if (is_empty(sections)):
    []
  else:
    do
      let section_contains = fn(section):
              if (_is_section(section)):
                do
                  section[:header] | to_text() | contains(text)
                end
              else:
                false
            end
      | filter(sections, section_contains)
    end
end

# Filters sections by a pattern match in the title text.
def title_match(sections, pattern):
  if (is_empty(sections)):
    []
  else:
    do
      let section_filter = fn(section):
              if (_is_section(section)):
                do
                  section[:header] | to_text() | regex_match(pattern) | !is_empty()
                end
              else:
                false
            end
      | filter(sections, section_filter)
    end
end

# Returns the title text of a section (header text without the # symbols).
def title(section):
  if (_is_section(section)):
    do section[:header] | to_text();
  else:
    ""
end

# Returns the content of a section (all nodes except the header).
# Deprecated: Use body() instead, as content()
def content(section):
  body(section)
end

# Returns the body of a section (all nodes except the header).
def body(section):
  if (_is_section(section)):
    section[:children]
  else:
    []
end

# Returns all nodes of a section, including both the header and content.
def all_nodes(section):
  if (_is_section(section)):
    do
      section[:header] + section[:children]
    end
  else:
    []
end

# Filters sections by heading level. l can be a number (exact level) or a range array (e.g. 1..2).
def by_level(sections, l):
  if (is_empty(sections)):
    []
  else:
    filter(sections, fn(section):
      if (_is_section(section)):
        do
          let lvl = level(section)
          | if (is_array(l)):
              contains(l, lvl)
            else:
              lvl == l
        end
      else:
        false
    end)
end

# Returns the header level (1-6) of a section.
def level(section):
  if (_is_section(section)):
    do
      let header = section[:header]
      | if (is_h_level(header, 1)): 1
        elif (is_h_level(header, 2)): 2
        elif (is_h_level(header, 3)): 3
        elif (is_h_level(header, 4)): 4
        elif (is_h_level(header, 5)): 5
        elif (is_h_level(header, 6)): 6
        else: 0
    end
  else:
    0
end

# Returns the nth section from an array of sections (0-indexed).
def nth(sections, n):
  if (is_empty(sections) || n < 0 || n >= len(sections)):
    None
  else:
    get(sections, n)
end

# Extracts titles from all sections.
def titles(sections):
  if (is_empty(sections)):
    []
  else:
    map(sections, title)
end

# Extracts body from all sections.
def bodies(sections):
  if (is_empty(sections)):
    []
  else:
    map(sections, body)
end

# Generates a table of contents from sections.
def toc(sections):
  if (is_empty(sections)):
    []
  else:
    do
      let create_toc_entry = fn(section):
              if (_is_section(section)):
                do
                  let lvl = level(section)
                  | let indent = join(map(range(0, lvl - 1), fn(_): "  ";), "")
                  | let title_text = title(section)
                  | indent + "- " + title_text
                end
              else:
                ""
            end
      | map(sections, create_toc_entry)
    end
end

# Checks if a section has any content beyond the header.
def has_content(section):
  if (!is_dict(section)):
    error("Expected a dictionary, but got a different type.")
  elif (_is_section(section)):
    len(content(section)) > 0
  else:
    false
end

# Flattens sections back to markdown nodes for output.
# This converts section objects back to their original markdown node arrays.
def collect(sections):
  if (is_empty(sections)):
    []
  else:
    flat_map(sections, fn(section):
      if (_is_section(section)):
        all_nodes(section)
      else:
        section
    end)
end
