native_elixir_pdf_utilities

defmodule NativeElixirPdfUtilities.Merge do
  @moduledoc """
  PDF utilities for merging documents through the shared native reader.

  Notes and constraints:
  - Emits a classic PDF 1.7 header and builds a fresh `xref` + `trailer`.
  - Resolves classic xref tables, xref streams, and object streams before copying
    active objects with fresh identifiers.
  - Adjusts indirect references (`n g R`) to the new numbering.
  - Collects Page objects and builds a new `Catalog` + `Pages` tree that references them.
  - Leaves stream bytes untouched and preserves declared `/Length` (direct or indirect ref),
    only renumbering indirect references as needed.

  The merger is conservative and pragmatic, targeting structural correctness for common PDFs.
  """

  alias NativeElixirPdfUtilities.Diagnostics
  alias NativeElixirPdfUtilities.Pdf.Reader
  alias NativeElixirPdfUtilities.Tokenizer

  @type pdf_bin :: binary()
  @typedoc "A single token as produced by `NativeElixirPdfUtilities.Tokenizer`."
  @type token :: Tokenizer.token()
  @typedoc "A list of PDF tokens."
  @type tokens :: [token()]
  @typedoc "Object record captured while indexing inputs."
  @type obj_rec :: %{obj: integer(), gen: integer(), tokens: tokens()}
  @typedoc "Mapping from original object id to new object id."
  @type id_map :: %{optional(integer()) => integer()}
  @typedoc "Byte-offset table for xref: object id -> {byte_offset, generation}."
  @type offsets_map :: %{optional(integer()) => {non_neg_integer(), non_neg_integer()}}
  @type error_reason :: :empty_pdf_list | :invalid_pdf_input

  @doc """
  Merge a list of PDF binaries into a single PDF binary.

  It resolves active objects through the shared reader, renumbers them, collects
  Page objects, and emits a new Catalog/Pages tree referencing all input pages.
  """
  @spec merge([pdf_bin()]) ::
          {:ok, pdf_bin()} | {:error, {error_reason(), Diagnostics.diagnostic()}}
  def merge(bins) do
    case bins do
      [] ->
        Diagnostics.error(:merge, :empty_pdf_list, "merge/1 expects at least one PDF binary",
          operation: :merge,
          module: __MODULE__
        )

      bins when is_list(bins) ->
        case Enum.all?(bins, &is_binary/1) do
          true ->
            do_merge(bins)

          false ->
            Diagnostics.error(
              :merge,
              :invalid_pdf_input,
              "merge/1 expects a list of PDF binaries",
              operation: :merge,
              module: __MODULE__
            )
        end

      _ ->
        Diagnostics.error(:merge, :invalid_pdf_input, "merge/1 expects a list of PDF binaries",
          operation: :merge,
          module: __MODULE__
        )
    end
  end

  defp do_merge(bins) do
    case Enum.reduce_while(bins, {:ok, []}, fn bin, {:ok, inputs} ->
           case index_pdf(bin) do
             {:ok, input} -> {:cont, {:ok, [input | inputs]}}
             {:error, {_reason, diagnostic}} -> {:halt, {:error, diagnostic}}
           end
         end) do
      {:ok, inputs} ->
        build_merged_pdf(Enum.reverse(inputs))

      {:error, diagnostic} ->
        Diagnostics.error(
          :merge,
          :invalid_pdf_input,
          "merge/1 received an invalid PDF: #{diagnostic.message}",
          operation: :merge,
          module: __MODULE__
        )
    end
  end

  defp build_merged_pdf(inputs) do
    # 2) Assign id offsets so object ids won't collide; reserve 1,2 for Pages,Catalog
    {inputs2, _next_id} = assign_offsets(inputs, 3)

    # 3) Collect all page ids in new numbering (flatten Pages)
    page_ids =
      inputs2
      |> Enum.flat_map(fn %{pages: pages, map: map} ->
        Enum.map(pages, fn {object, generation} ->
          {Map.fetch!(map, object), generation}
        end)
      end)
      |> Enum.reduce({[], MapSet.new()}, fn id, {acc, seen} ->
        if MapSet.member?(seen, id), do: {acc, seen}, else: {[id | acc], MapSet.put(seen, id)}
      end)
      |> then(fn {acc, _} -> Enum.reverse(acc) end)

    pages_obj_id = 1
    catalog_obj_id = 2

    # 4) Render all objects with rewritten refs
    # We'll render: new Pages, new Catalog, then all rewritten input objects
    {pieces, offsets, pos} = add_piece([], pdf_header(), %{}, 0)
    render_pages = render_pages_object(pages_obj_id, page_ids)
    {pieces, offsets, pos} = add_object(pieces, offsets, pos, pages_obj_id, 0, render_pages)
    render_catalog = render_catalog_object(catalog_obj_id, pages_obj_id)

    {pieces, offsets, pos} =
      add_object(pieces, offsets, pos, catalog_obj_id, 0, render_catalog)

    {pieces, offsets, pos} =
      Enum.reduce(inputs2, {pieces, offsets, pos}, fn input = %{objects: objs, map: map}, acc ->
        Enum.reduce(objs, acc, fn obj, {pieces, offsets, pos} ->
          new_id = Map.fetch!(map, obj.obj)
          page_ctx = page_injection_ctx(obj, input, pages_obj_id)
          body = render_object_body(obj.tokens, map, page_ctx)
          add_object(pieces, offsets, pos, new_id, obj.gen, body)
        end)
      end)

    # 5) Xref + trailer
    max_obj_id = Enum.max([catalog_obj_id, pages_obj_id | Map.keys(offsets)])
    {xref_io, _xref_pos} = xref_and_trailer(offsets, pos, max_obj_id, catalog_obj_id)

    final_io = [Enum.reverse(pieces), xref_io]
    {:ok, IO.iodata_to_binary(final_io)}
  end

  # === Indexing ===

  # Index a PDF binary into objects, page ids and inherited attributes.
  defp index_pdf(bin) do
    case Reader.read(bin) do
      {:ok, document} ->
        objects =
          document.objects
          |> Enum.reject(fn {_ref, object} -> structural_reader_object?(object.value) end)
          |> Enum.map(fn {{object, generation}, parsed} ->
            %{obj: object, gen: generation, tokens: parsed.tokens}
          end)
          |> Enum.sort_by(&{&1.obj, &1.gen})

        root_pages = find_root_pages_id(objects)

        {:ok,
         %{
           objects: objects,
           pages: Enum.map(document.pages, & &1.ref),
           root_pages: root_pages,
           inherited: page_inheritances(objects, root_pages),
           max_obj: Enum.reduce(objects, 0, fn object, maximum -> max(maximum, object.obj) end)
         }}

      {:error, _} = reader_error ->
        reader_error
    end
  end

  defp structural_reader_object?(value) do
    case value do
      %{"Type" => {:name, type}} when type in ["XRef", "ObjStm"] -> true
      _ -> false
    end
  end

  # Best-effort predicate to detect a Page dictionary by scanning for /Type /Page.
  defp object_is_page?(tokens) do
    # Best-effort: look for /Type /Page anywhere
    Enum.chunk_every(tokens, 2, 1, :discard)
    |> Enum.any?(fn
      [{:name, "Type"}, {:name, "Page"}] -> true
      _ -> false
    end)
  end

  # Find the object id of the root Pages tree via the Catalog's /Pages reference.
  defp find_root_pages_id(objects) do
    %{tokens: tokens} =
      Enum.find(objects, fn %{tokens: tokens} ->
        Enum.chunk_every(tokens, 2, 1, :discard)
        |> Enum.any?(fn
          [{:name, "Type"}, {:name, "Catalog"}] -> true
          _ -> false
        end)
      end)

    find_pages_ref_in_tokens(tokens)
  end

  # Look for '/Pages <obj> <gen> R' in a token sequence and return <obj>.
  defp find_pages_ref_in_tokens(tokens) do
    tokens
    |> Enum.chunk_every(4, 1, :discard)
    |> Enum.find_value(fn
      [{:name, "Pages"}, {:int, obj}, {:int, _gen}, :R] -> obj
      _ -> nil
    end)
  end

  # Resolve inheritable page attributes through the complete /Pages tree.
  defp page_inheritances(objects, root_pages_id) do
    object_by_id = Map.new(objects, &{&1.obj, &1})

    collect_page_inheritances(
      root_pages_id,
      object_by_id,
      %{resources: nil, mediabox: nil},
      %{}
    )
  end

  defp collect_page_inheritances(object_id, object_by_id, inherited, page_attributes) do
    %{tokens: tokens} = object = Map.fetch!(object_by_id, object_id)

    attributes = %{
      resources: find_value_after_name(tokens, "Resources") || inherited.resources,
      mediabox: find_value_after_name(tokens, "MediaBox") || inherited.mediabox
    }

    case object_is_page?(tokens) do
      true ->
        Map.put(page_attributes, object.obj, attributes)

      false ->
        page_child_references(tokens)
        |> Enum.reduce(page_attributes, fn child_id, acc ->
          collect_page_inheritances(child_id, object_by_id, attributes, acc)
        end)
    end
  end

  defp page_child_references(tokens) do
    [:lbracket | children] = find_value_after_name(tokens, "Kids")

    children
    |> Enum.drop(-1)
    |> Enum.chunk_every(3)
    |> Enum.map(fn [{:int, object_id}, {:int, _generation}, :R] -> object_id end)
  end

  # Find the value tokens immediately following a given name key in a token list.
  defp find_value_after_name(tokens, name) do
    case Enum.split_while(tokens, fn t -> t != {:name, name} end) do
      {_prefix, []} ->
        nil

      {_prefix, [_name | rest]} ->
        {value, _rest} = read_value_tokens(rest)
        value
    end
  end

  # Read a single value (dict/array/ref/atom) from a token stream and also return the rest.
  defp read_value_tokens(tokens) do
    case tokens do
      [:dict_start | rest] ->
        take_until_matching(rest, :dict_start, :dict_end, 1, [:dict_start])

      [:lbracket | rest] ->
        take_until_matching(rest, :lbracket, :rbracket, 1, [:lbracket])

      [{:int, a}, {:int, b}, :R | rest] ->
        {[{:int, a}, {:int, b}, :R], rest}

      [tok | rest] ->
        {[tok], rest}
    end
  end

  # Collect a balanced dictionary or array value while preserving nested tokens.
  defp take_until_matching(tokens, opening, closing, depth, acc) do
    case tokens do
      [^opening | rest] ->
        take_until_matching(rest, opening, closing, depth + 1, [opening | acc])

      [^closing | rest] when depth == 1 ->
        {Enum.reverse([closing | acc]), rest}

      [^closing | rest] ->
        take_until_matching(rest, opening, closing, depth - 1, [closing | acc])

      [token | rest] ->
        take_until_matching(rest, opening, closing, depth, [token | acc])
    end
  end

  # Build a page-rewrite context for Page objects (to set Parent/Resources/MediaBox), else nil.
  defp page_injection_ctx(object, %{inherited: inheritances}, parent_id) do
    if object_is_page?(object.tokens) do
      inherited = Map.get(inheritances, object.obj, %{resources: nil, mediabox: nil})

      %{
        parent_id: parent_id,
        resources_tokens: inherited.resources,
        mediabox_tokens: inherited.mediabox
      }
    else
      nil
    end
  end

  # Assign id ranges to each input to avoid collisions; returns inputs augmented with :map.
  defp assign_offsets(inputs, start_id) do
    Enum.map_reduce(inputs, start_id, fn %{max_obj: max_obj} = input, next_id ->
      base = next_id
      map = fn orig -> base + orig end
      # Build id map for all ids 1..max_obj that actually appear
      ids_present = MapSet.new(Enum.map(input.objects, & &1.obj))

      id_map =
        ids_present
        |> Enum.into(%{}, fn id -> {id, map.(id)} end)

      {Map.put(input, :map, id_map), base + max_obj + 1}
    end)
  end

  # === Rendering ===

  # Return a fixed classic PDF header with a binary-comment line.
  defp pdf_header do
    ["%PDF-1.7\n%\xE2\xE3\xCF\xD3\n"]
  end

  # Append a chunk to the output pieces and update the running byte position.
  defp add_piece(pieces, piece, offsets, pos) do
    len = :erlang.iolist_size(piece)
    {[piece | pieces], offsets, pos + len}
  end

  # Append a fully formatted object to the output, recording its starting offset.
  defp add_object(pieces, offsets, pos, id, gen, body_io) do
    header = [Integer.to_string(id), " ", Integer.to_string(gen), " obj\n"]
    footer = "\nendobj\n"
    piece = [header, body_io, footer]
    len = :erlang.iolist_size(piece)
    offsets2 = Map.put(offsets, id, {pos, gen})
    {[piece | pieces], offsets2, pos + len}
  end

  # Render the top-level Pages dictionary referencing all collected Page kids.
  defp render_pages_object(_pages_obj_id, page_ids) do
    kids_refs =
      page_ids
      |> Enum.map(fn {id, generation} ->
        [Integer.to_string(id), " ", Integer.to_string(generation), " R"]
      end)
      |> Enum.intersperse(" ")

    [
      "<< /Type /Pages /Kids [ ",
      kids_refs,
      " ] /Count ",
      Integer.to_string(length(page_ids)),
      " >>\n"
    ]
  end

  # Render the top-level Catalog referencing the generated Pages object.
  defp render_catalog_object(_catalog_obj_id, pages_obj_id) do
    [
      "<< /Type /Catalog /Pages ",
      Integer.to_string(pages_obj_id),
      " 0 R >>\n"
    ]
  end

  # Render an object's body while optionally rewriting Page dict content and remapping refs.
  defp render_object_body(tokens, id_map, page_ctx) do
    # If this is a Page dict, rewrite Parent and ensure Resources/MediaBox
    tokens2 =
      case page_ctx do
        nil -> tokens
        %{parent_id: _} -> rewrite_page_tokens(tokens, page_ctx)
      end

    # Replace indirect references with mapped ids; render tokens with spaces
    render_tokens(tokens2, id_map)
  end

  # Rewrite top-level Page dictionary to set Parent, ensure /Type /Page, /Resources and /MediaBox.
  defp rewrite_page_tokens(tokens, %{
         parent_id: parent_id,
         resources_tokens: inh_res,
         mediabox_tokens: inh_mb
       }) do
    # We expect a single top-level dict in a Page object. Split it out, sanitize, and put it back.
    [:dict_start | rest] = tokens
    {dict_inner, before, afterr} = do_take_dict(rest, 1, [], [])

    dict_inner =
      dict_inner
      |> drop_key("Parent")
      |> put_key("Parent", [{:generated_reference, parent_id}])
      |> ensure_type_page()
      |> ensure_resources(inh_res)
      |> ensure_mediabox(inh_mb || default_mediabox())

    before ++ [:dict_start | dict_inner] ++ [:dict_end | afterr]
  end

  # Worker for top-level dict extraction.
  defp do_take_dict(tokens, depth, acc, before) do
    case tokens do
      [:dict_start | rest] ->
        do_take_dict(rest, depth + 1, [:dict_start | acc], before)

      [:dict_end | rest] when depth == 1 ->
        {Enum.reverse(acc), Enum.reverse(before), rest}

      [:dict_end | rest] ->
        do_take_dict(rest, depth - 1, [:dict_end | acc], before)

      [token | rest] ->
        do_take_dict(rest, depth, [token | acc], before)
    end
  end

  # Drop a key (and its value) from a flat dict token list if present.
  defp drop_key(tokens, name) do
    case split_on_name(tokens, name) do
      {:ok, left, _val, right} -> left ++ right
      :error -> tokens
    end
  end

  # Put or replace a key with the given value tokens, appending near the end by default.
  defp put_key(tokens, name, value_tokens) do
    # Put/replace near the end so it’s visible
    case split_on_name(tokens, name) do
      {:ok, left, _old, right} -> left ++ [{:name, name} | value_tokens] ++ right
      :error -> tokens ++ [{:name, name} | value_tokens]
    end
  end

  # Ensure /Type /Page is set.
  defp ensure_type_page(tokens) do
    put_key(tokens, "Type", [{:name, "Page"}])
  end

  # Keep existing /Resources if non-empty; otherwise inject inherited /Resources when available.
  defp ensure_resources(tokens, inh_res) do
    case split_on_name(tokens, "Resources") do
      {:ok, left, val, right} when is_list(val) and val != [] ->
        left ++ [{:name, "Resources"} | val] ++ right

      _ ->
        if is_list(inh_res) and inh_res != [] do
          put_key(tokens, "Resources", inh_res)
        else
          tokens
        end
    end
  end

  # Keep a valid /MediaBox array; otherwise use provided fallback.
  defp ensure_mediabox(tokens, fallback) do
    fallback = if valid_box?(fallback), do: fallback, else: default_mediabox()

    case split_on_name(tokens, "MediaBox") do
      {:ok, left, val, right} ->
        if valid_box?(val) do
          left ++ [{:name, "MediaBox"} | val] ++ right
        else
          left ++ [{:name, "MediaBox"} | fallback] ++ right
        end

      :error ->
        put_key(tokens, "MediaBox", fallback)
    end
  end

  # A4
  defp default_mediabox do
    [:lbracket, {:int, 0}, {:int, 0}, {:int, 595}, {:int, 842}, :rbracket]
  end

  # Validate a box array [a b c d] of numbers.
  defp valid_box?(tokens) do
    case tokens do
      [:lbracket, a, b, c, d, :rbracket] ->
        Enum.all?([a, b, c, d], &numeric_token?/1)

      _ ->
        false
    end
  end

  defp numeric_token?(token) do
    case token do
      {:int, _value} -> true
      {:real, _value} -> true
      _ -> false
    end
  end

  # Split a flat dict token list on a name key; returns left, value tokens, right or :error.
  defp split_on_name(tokens, name) do
    # returns {:ok, left, value_tokens, right} or :error
    case Enum.split_while(tokens, fn t -> t != {:name, name} end) do
      {_left, []} ->
        :error

      {left, [_name | rest]} ->
        {val, rest2} = read_value_tokens(rest)
        {:ok, left, val, rest2}
    end
  end

  # Render tokens back into iodata while remapping indirect references using id_map.
  defp render_tokens(tokens, id_map) do
    do_render_tokens(tokens, id_map, [], nil) |> Enum.reverse()
  end

  defp do_render_tokens(tokens, id_map, acc, _last_name) do
    case tokens do
      [] ->
        acc

      [{:name, name} | rest] ->
        do_render_tokens(rest, id_map, [["/", name] | add_sep(acc)], name)

      [{:generated_reference, obj} | rest] ->
        io = [Integer.to_string(obj), " 0 R"]
        do_render_tokens(rest, id_map, [io | add_sep(acc)], nil)

      [{:int, obj}, {:int, gen}, :R | rest] ->
        new_obj = Map.get(id_map, obj, obj)
        io = [Integer.to_string(new_obj), " ", Integer.to_string(gen), " R"]
        do_render_tokens(rest, id_map, [io | add_sep(acc)], nil)

      [:dict_start | rest] ->
        do_render_tokens(rest, id_map, ["<<" | add_sep(acc)], nil)

      [:dict_end | rest] ->
        do_render_tokens(rest, id_map, [">>" | add_sep(acc)], nil)

      [:lbracket | rest] ->
        do_render_tokens(rest, id_map, ["[" | add_sep(acc)], nil)

      [:rbracket | rest] ->
        do_render_tokens(rest, id_map, ["]" | add_sep(acc)], nil)

      [:stream, {:stream_data, data}, :endstream | rest] ->
        do_render_tokens(rest, id_map, [["\nstream\n", data, "\nendstream"] | acc], nil)

      [{:string, string} | rest] ->
        do_render_tokens(rest, id_map, [["(", escape_literal(string), ")"] | add_sep(acc)], nil)

      [{:hex_string, string} | rest] ->
        do_render_tokens(rest, id_map, [["<", to_hex(string), ">"] | add_sep(acc)], nil)

      [{:int, int} | rest] ->
        do_render_tokens(rest, id_map, [Integer.to_string(int) | add_sep(acc)], nil)

      [{:real, real} | rest] ->
        do_render_tokens(rest, id_map, [format_pdf_real(real) | add_sep(acc)], nil)

      [true | rest] ->
        do_render_tokens(rest, id_map, ["true" | add_sep(acc)], nil)

      [false | rest] ->
        do_render_tokens(rest, id_map, ["false" | add_sep(acc)], nil)

      [:null | rest] ->
        do_render_tokens(rest, id_map, ["null" | add_sep(acc)], nil)
    end
  end

  # Add a separating space in the output unless at the beginning.
  defp add_sep(acc) do
    case acc do
      [] -> []
      _ -> [" " | acc]
    end
  end

  # Ensure reals are rendered as plain decimal (no scientific notation)
  defp format_pdf_real(f) do
    # If the value is essentially an integer, emit as integer
    i = trunc(f)

    if abs(f - i) < 1.0e-9 do
      Integer.to_string(i)
    else
      # Fixed decimals with trimming
      s = :erlang.float_to_binary(f, [{:decimals, 10}])
      trim_trailing_zeros_and_dot(s)
    end
  end

  defp trim_trailing_zeros_and_dot(bin) do
    bin
    |> String.replace_trailing("0", "")
    |> String.replace_trailing(".", "")
  end

  # Escape a literal string for inclusion in (...) with PDF-compliant escapes.
  defp escape_literal(bin) do
    bin
    |> :binary.bin_to_list()
    |> Enum.map(fn
      ?\n ->
        "\\n"

      ?\r ->
        "\\r"

      ?\t ->
        "\\t"

      ?\b ->
        "\\b"

      ?\f ->
        "\\f"

      ?( ->
        "\\("

      ?) ->
        "\\)"

      ?\\ ->
        "\\\\"

      c when c < 32 or c > 126 ->
        # Use octal escape for non-printable
        :io_lib.format("\\~.3.0b", [c])

      c ->
        <<c>>
    end)
  end

  # Convert bytes to uppercase hex pairs iodata.
  defp to_hex(bin) do
    for <<c <- bin>>, into: [], do: :io_lib.format("~2.16.0B", [c])
  end

  # Build a classic xref table and trailer for the accumulated object offsets.
  defp xref_and_trailer(offsets, pos, max_id, root_id) do
    xref_pos = pos
    size = max_id + 1

    header = ["xref\n0 ", Integer.to_string(size), "\n"]

    all_ids = Enum.to_list(0..max_id)

    # Free ids are those without offsets (0 is always free)
    nonzero_free =
      all_ids
      |> Enum.reject(&(&1 == 0))
      |> Enum.filter(&(not Map.has_key?(offsets, &1)))

    # Build free-list mapping for nonzero free objects: id -> next_id (last points to 0)
    next_of =
      nonzero_free
      |> Enum.zip(Enum.drop(nonzero_free, 1) ++ [0])
      |> Map.new()

    # Object 0 must point to the first free object (or 0 if none)
    first_free = List.first(nonzero_free) || 0

    entries =
      Enum.map(all_ids, fn id ->
        case Map.fetch(offsets, id) do
          {:ok, {off, gen}} ->
            [pad10(off), " ", pad5(gen), " n \n"]

          :error when id == 0 ->
            [pad10(first_free), " 65535 f \n"]

          :error ->
            next = Map.get(next_of, id, 0)
            [pad10(next), " 00000 f \n"]
        end
      end)

    trailer = [
      "trailer\n<< /Size ",
      Integer.to_string(size),
      " /Root ",
      Integer.to_string(root_id),
      " 0 R >>\n",
      "startxref\n",
      Integer.to_string(xref_pos),
      "\n%%EOF\n"
    ]

    {[header, entries, trailer], xref_pos}
  end

  # Pad an integer to 10 digits with leading zeroes.
  defp pad10(int) do
    s = Integer.to_string(int)
    pad = 10 - byte_size(s)
    if pad > 0, do: :binary.copy("0", pad) <> s, else: s
  end

  # Pad an integer to 5 digits with leading zeroes.
  defp pad5(int) do
    s = Integer.to_string(int)
    pad = 5 - byte_size(s)
    if pad > 0, do: :binary.copy("0", pad) <> s, else: s
  end
end