lang_extract

defmodule LangExtract.Chunker do
  @moduledoc """
  Splits text into sentence-level chunks using the Alignment.Tokenizer.

  Sentence boundary rules (mirroring upstream's `find_sentence_range`):
  1. A `:punctuation` token ending in a sentence terminator (`.`, `!`, `?`,
     CJK equivalents — `...` counts, since symbol runs are one token) ends a
     sentence, unless the previous token plus the terminator form a known
     abbreviation (`"Dr" <> "." == "Dr."`).
  2. After sentence-ending punctuation, trailing closing punctuation
     (`"`, `'`, `)`, `]`, `}`, `»`, `\u201D`, `\u2019`) is consumed into the same sentence.
  3. A `:whitespace` token containing `\\n` starts a new sentence unless the
     next token begins lowercase — lines opening with quotes, digits, or
     capitals all break (upstream: "assume break unless lowercase").
  """

  alias LangExtract.Alignment.Tokenizer
  alias LangExtract.Chunker.Chunk

  @abbreviations ~w(Mr. Mrs. Ms. Dr. Prof. St.)
  @closing_punctuation [~s("), "'", ")", "]", "}", "»", "\u201D", "\u2019"]
  # Upstream's _END_OF_SENTENCE_PATTERN: a token ending in a sentence
  # terminator (same-symbol runs make "..." one token, so match the tail).
  @sentence_ending ~r/[.?!。！？\x{0964}]["'”’»)\]}]*$/u

  @doc """
  Splits text into chunks respecting sentence boundaries.

  ## Options

    * `:max_chunk_chars` — maximum characters per chunk (required).
      Char-denominated to mirror upstream's `max_char_buffer`, so chunk
      boundaries land identically across the two libraries — the
      cross-library benchmarks depend on that. Output offsets are bytes.

  """
  @spec chunk(String.t(), keyword()) :: [Chunk.t()]
  def chunk(text, opts) when is_binary(text) do
    max_chars = Keyword.fetch!(opts, :max_chunk_chars)

    text
    |> find_sentences()
    |> Enum.flat_map(&split_oversized(&1, max_chars))
    |> pack_sentences(max_chars)
  end

  # A sentence longer than the budget is hard-split at token boundaries
  # (mirroring upstream's ChunkIterator) — otherwise boundary-free text
  # (logs, minified content) becomes one whole-document chunk and defeats
  # the size guarantee. Fragments concatenate back to the sentence exactly,
  # so pack_sentences' byte accounting is unaffected. A single token longer
  # than the budget stays whole: token boundaries are never violated.
  defp split_oversized(sentence, max_chars) do
    # chars ≤ bytes, so a small byte count proves the sentence fits without
    # walking its graphemes; only near-oversized sentences pay for a count.
    if byte_size(sentence) <= max_chars or String.length(sentence) <= max_chars do
      [sentence]
    else
      sentence
      |> Tokenizer.tokenize()
      |> split_tokens(max_chars)
    end
  end

  defp split_tokens(tokens, max_chars) do
    {fragments, current, _len} =
      Enum.reduce(tokens, {[], [], 0}, fn token, {fragments, current, len} ->
        token_len = String.length(token.text)

        cond do
          current == [] ->
            {fragments, [token.text], token_len}

          len + token_len <= max_chars ->
            {fragments, [token.text | current], len + token_len}

          true ->
            {[join_fragment(current) | fragments], [token.text], token_len}
        end
      end)

    Enum.reverse([join_fragment(current) | fragments])
  end

  defp join_fragment(reversed_texts) do
    reversed_texts |> Enum.reverse() |> IO.iodata_to_binary()
  end

  defp pack_sentences(sentences, max_chars) do
    {chunks, current_text, current_start, _current_len} =
      Enum.reduce(sentences, {[], "", 0, 0}, fn sentence,
                                                {chunks, current_text, current_start, current_len} ->
        sentence_len = String.length(sentence)

        if current_len + sentence_len <= max_chars or current_text == "" do
          {chunks, current_text <> sentence, current_start, current_len + sentence_len}
        else
          byte_end = current_start + byte_size(current_text)
          chunk = %Chunk{text: current_text, byte_start: current_start, byte_end: byte_end}
          {[chunk | chunks], sentence, byte_end, sentence_len}
        end
      end)

    if current_text != "" do
      byte_end = current_start + byte_size(current_text)

      Enum.reverse([
        %Chunk{text: current_text, byte_start: current_start, byte_end: byte_end} | chunks
      ])
    else
      Enum.reverse(chunks)
    end
  end

  # Public only as a test seam: the sentence-boundary rules aren't
  # observable through chunk/2 (packing merges sentences back together).
  @doc false
  @spec find_sentences(String.t()) :: [String.t()]
  def find_sentences(""), do: []

  def find_sentences(text) when is_binary(text) do
    tokens = Tokenizer.tokenize(text)
    tokens_tuple = List.to_tuple(tokens)
    count = tuple_size(tokens_tuple)
    boundaries = find_boundaries(tokens_tuple, count)
    tokens_to_sentences(tokens_tuple, boundaries, text)
  end

  defp find_boundaries(tokens_tuple, count) do
    boundaries =
      Enum.reduce(0..(count - 1)//1, [], fn idx, acc ->
        token = elem(tokens_tuple, idx)

        cond do
          sentence_end_by_punctuation?(token, idx, tokens_tuple) ->
            end_idx = consume_closing_punctuation(idx + 1, tokens_tuple, count)
            [end_idx | acc]

          sentence_end_by_newline?(token, idx, tokens_tuple, count) ->
            [idx | acc]

          true ->
            acc
        end
      end)

    boundaries
    |> Enum.sort()
    |> Enum.uniq()
    |> ensure_final_boundary(count)
  end

  defp sentence_end_by_punctuation?(%{type: :punctuation, text: text}, idx, tokens_tuple) do
    Regex.match?(@sentence_ending, text) and not abbreviation_before?(idx, tokens_tuple, text)
  end

  defp sentence_end_by_punctuation?(_token, _idx, _tokens_tuple), do: false

  # Upstream concatenates the previous token with the terminator and
  # checks the pair ("Dr" <> "." == "Dr."), so "Dr..." still breaks.
  defp abbreviation_before?(punct_idx, tokens_tuple, punct_text) when punct_idx > 0 do
    prev = elem(tokens_tuple, punct_idx - 1)
    (prev.text <> punct_text) in @abbreviations
  end

  defp abbreviation_before?(_punct_idx, _tokens_tuple, _punct_text), do: false

  defp consume_closing_punctuation(idx, tokens_tuple, count) when idx < count do
    token = elem(tokens_tuple, idx)

    if token.type == :punctuation and token.text in @closing_punctuation do
      consume_closing_punctuation(idx + 1, tokens_tuple, count)
    else
      idx
    end
  end

  defp consume_closing_punctuation(idx, _tokens_tuple, _count), do: idx

  # Upstream: "Assume break unless lowercase (covers numbers/quotes)" —
  # a line starting with “, a digit, or an uppercase letter all break.
  defp sentence_end_by_newline?(%{type: :whitespace, text: ws_text}, idx, tokens_tuple, count) do
    if String.contains?(ws_text, "\n") and idx + 1 < count do
      next = elem(tokens_tuple, idx + 1)
      not lowercase_start?(next.text)
    else
      false
    end
  end

  defp sentence_end_by_newline?(_token, _idx, _tokens_tuple, _count), do: false

  defp lowercase_start?(<<first::utf8, _rest::binary>>) do
    char = <<first::utf8>>
    String.downcase(char) == char and String.upcase(char) != char
  end

  defp lowercase_start?(_), do: false

  defp ensure_final_boundary(boundaries, count) do
    if List.last(boundaries) == count do
      boundaries
    else
      boundaries ++ [count]
    end
  end

  # Extracts sentence strings using byte offsets from the token tuple.
  # Uses boundary pairs to look up first/last token directly — O(1) per sentence.
  defp tokens_to_sentences(tokens_tuple, boundaries, text) do
    {sentences, _} =
      Enum.reduce(boundaries, {[], 0}, fn boundary, {acc, start_idx} ->
        if boundary > start_idx do
          first = elem(tokens_tuple, start_idx)
          last = elem(tokens_tuple, boundary - 1)
          sentence = binary_part(text, first.byte_start, last.byte_end - first.byte_start)
          {[sentence | acc], boundary}
        else
          {acc, boundary}
        end
      end)

    Enum.reverse(sentences)
  end
end