chunkex

defmodule Chunkex do
  @moduledoc """
  A powerful Elixir tool for intelligently chunking your project's source code
  to support Local RAG (Retrieval-Augmented Generation) systems.

  Chunkex extracts functions from your Elixir codebase and creates structured
  chunks optimized for embedding and retrieval. This enables AI agents to
  access relevant code snippets with high precision, improving response quality
  while reducing token usage.

  ## Usage

      # Chunk the current directory
      Chunkex.chunk()

      # Chunk a specific directory
      Chunkex.chunk("/path/to/your/project")

  ## Output

  Generates a `tmp/chunks.jsonl` file containing structured chunks with metadata
  including file paths, line ranges, language identification, and SHA hashes.
  """

  @doc """
  Chunks all Elixir files in the specified directory and writes the results
  to `tmp/chunks.jsonl`.

  ## Parameters

  * `root` - The root directory to chunk (defaults to current directory)

  ## Returns

  Returns `:ok` and prints the number of chunks generated.

  ## Examples

      Chunkex.chunk()
      #=> Generated 15 chunks in ./tmp/chunks.jsonl

      Chunkex.chunk("/path/to/project")
      #=> Generated 23 chunks in /path/to/project/tmp/chunks.jsonl

  """
  def chunk(root \\ ".") do
    files = Path.wildcard(Path.join([root, "lib/**/*.ex"]))

    chunks = Enum.flat_map(files, &chunk_file/1)

    # Ensure tmp directory exists
    tmp_dir = Path.join(root, "tmp")
    File.mkdir_p!(tmp_dir)

    # Write chunks to tmp/chunks.jsonl
    output_path = Path.join(tmp_dir, "chunks.jsonl")

    chunks
    |> Enum.map(&Jason.encode!/1)
    |> Enum.join("\n")
    |> then(&File.write!(output_path, &1))

    IO.puts("Generated #{length(chunks)} chunks in #{output_path}")
  end

  defp chunk_file(path) do
    src = File.read!(path)
    sha = :crypto.hash(:sha256, src) |> Base.encode16(case: :lower)

    try do
      {:ok, ast} = Sourceror.parse_string(src, token_metadata: true, columns: true)

      funs = ast
      |> Macro.prewalk([], fn
        {kind, meta, _} = node, acc when kind in [:def, :defp, :defmacro] ->
          {node, [extract(meta, path, src) | acc]}
        node, acc -> {node, acc}
      end)
      |> elem(1)
      |> Enum.reverse()

      funs |> Enum.reverse() |> Enum.map(&Map.put(&1, "sha", sha))
    rescue
      _ -> []
    end
  end

  defp extract(meta, path, src) do
    ls = meta[:line]
    le = case meta[:end] do
      %{line: end_line} -> end_line
      _ -> ls + 30
    end
    text = src |> String.split("\n") |> Enum.slice(ls - 1, le - ls + 1) |> Enum.join("\n")
    %{"path" => path, "line_start" => ls, "line_end" => le, "lang" => "elixir", "text" => text}
  rescue
    _ -> %{"path" => path, "line_start" => 1, "line_end" => 1, "lang" => "elixir", "text" => "error"}
  end
end