local write_pdf = require'luatagging-objectwriter'
local write_pdf_utils = require'luatagging-objectwriter-utils'

local record ParentTree
  add: function(write_pdf.PdfObject): integer
  write_root: function(): write_pdf.Ref
end
local parenttree: ParentTree = {}

local max_width = 10
local entry_count: integer = 0
local tree: {TreeEntry} = {}

local record TreeEntry is {write_pdf.PdfObject}
  first: integer
end

function parenttree.add(obj: write_pdf.PdfObject): integer
  local id = entry_count
  entry_count = entry_count + 1
  local index = id
  local depth = 1
  while true do
    local current_block = tree[depth]
    if current_block == nil then
      current_block = {
        first = 0
      }
      write_pdf.array(current_block) -- This makes TreeEntry a subtype of Array
      tree[depth] = current_block
    end
    local new_width = (index % max_width) + 1
    if depth == 1 then
      current_block[2 * new_width - 1] = id
      current_block[2 * new_width] = obj
    else
      current_block[new_width] = obj
    end
    if new_width ~= max_width then break end
    local n = write_pdf.dict {
      Limits = write_pdf.array { current_block.first, id },
      [depth == 1 and 'Nums' or 'Kids'] = current_block as write_pdf.Array, -- Safe since we made TreeEntry a Array subtype
    }
    obj = write_pdf_utils.indirect(n)
    for i = 1, (depth == 1 and 2 or 1) * max_width do
      current_block[i] = nil
    end
    current_block.first = entry_count
    depth = depth + 1
    index = index // max_width
  end
  return id
end

function parenttree.write_root(): write_pdf.Ref
  if entry_count == 0 then return nil end
  local index = entry_count
  entry_count = entry_count - 1
  local max_depth = #tree
  local obj: write_pdf.Ref
  for depth, current_block in ipairs(tree) do
    if obj ~= nil then
      current_block[(index % max_width) + 1] = obj
    end
    if #current_block ~= 0 then
      local n = write_pdf.dict {
        [depth == 1 and 'Nums' or 'Kids'] = current_block as write_pdf.Array, -- Safe since we made TreeEntry a Array subtype
      }
      if depth ~= max_depth then
        n.Limits = write_pdf.array { current_block.first, entry_count }
      end
      obj = write_pdf_utils.indirect(n as write_pdf.Array)
    end
    index = index // max_width
  end
  return obj
end

return parenttree