-- tagpdf allows a bunch of different identifiers to be used to refer to structure elements,
-- mainly
--  - (structure) numbers: Internal ID available to each structure element.
--  - label: Label assigned for some elements.
--  - dest: LaTeX label strings which has been associated with a structure element.
--
-- In luatagging we unify labels and dests into towo namespaces for the same mechanism
-- and handle structure numbers similarly. This also means that there is no unique structure number
-- each structure element.

local structelem = require'luatagging-structelem'

local pending_metatable: metatable<structelem.PendingElement> = {
  __index = {
    resolve = function(self: structelem.PendingElement): structelem.StructureElement
      local resolved = self.resolved
      if not resolved then
        tex.error("Unresolved label used")
      end
      return resolved
    end,
  },
}

local record PendingParentElement is structelem.ReferencedElement
  base: structelem.ReferencedElement
end

local pending_parent_metatable: metatable<PendingParentElement> = {
  __index = {
    resolve = function(self: PendingParentElement): structelem.StructureElement
      local base = self.base:resolve()
      if not base then
        return
      end
      return base.parent as structelem.StructureElement
    end,
  },
}

local registry: {string: {string: structelem.ReferencedElement}} = {
  label = {},
  dest = {},
  struct_num = {},
}
local current_num: integer

local function get_labeled(labelKind: string, label: string): structelem.ReferencedElement
  local kindRegistry = registry[labelKind]
  if not kindRegistry then
    error'Invalid registry'
  end
  local element = kindRegistry[label]
  if not element then
    local pending: structelem.PendingElement = setmetatable({}, pending_metatable)
    element = pending
    kindRegistry[label] = element
  end
  return element
end

local function get_labeled_must_exist(labelKind: string, label: string): structelem.StructureElement
  local kindRegistry = registry[labelKind]
  if not kindRegistry then
    error'Invalid registry'
  end
  local element = kindRegistry[label]
  if element and getmetatable(element) ~= pending_metatable as metatable<structelem.ReferencedElement> then
    return element as structelem.StructureElement
  end
  return tex.error(string.format("Reference to unknown label %s of kind %s", label, labelKind))
end

local function exist(labelKind: string, label: string): boolean
  local kindRegistry = registry[labelKind]
  if not kindRegistry then
    error'Invalid registry'
  end
  local element = kindRegistry[label]
  return element and getmetatable(element) ~= pending_metatable as metatable<structelem.ReferencedElement> or false
end

local function assign_label(element: StructureElement, labelKind: string, label: string)
  local kindRegistry = registry[labelKind]
  if not kindRegistry then
    error'Invalid registry'
  end
  local existingReference = kindRegistry[label]
  if existingReference then
    if getmetatable(existingReference) == pending_metatable as metatable<structelem.ReferencedElement> then
      local pending = existingReference as structelem.PendingElement
      pending.resolved = element
    elseif existingReference == element then -- This happens surprisingly often for destinations...
      return
    else
      return tex.error(string.format("Conflicting assignments for label %s of kind  %s", label, labelKind))
    end
  end
  kindRegistry[label] = element
end

local function assign_num(element: StructureElement): integer
  local kindRegistry = registry.struct_num
  if current_num and kindRegistry[current_num as string] == element then
    return current_num
  end
  current_num = (current_num or 0) + 1
  kindRegistry[current_num as string] = element
  return current_num
end

local record Labels
  get_by_num: function(integer): structelem.StructureElement
  get: function(string, string): structelem.ReferencedElement
  get_must_exist: function(string, string): structelem.StructureElement
  assign: function(structelem.StructureElement, string, string)
  assign_num: function(structelem.StructureElement): integer
  get_parent: function(structelem.ReferencedElement): structelem.ReferencedElement
  exist_num: function(integer): boolean
  exist: function(string, string): boolean
end

local labels: Labels = {
  get_by_num = function(num: integer): structelem.StructureElement
    return get_labeled_must_exist('struct_num', num as string)
  end,
  get = get_labeled,
  get_must_exist = get_labeled_must_exist,
  assign = assign_label,
  assign_num = assign_num,
  get_parent = function(base: structelem.ReferencedElement): structelem.ReferencedElement
    local existingParent = (base as structelem.StructureElement).parent -- StructureElement is the only ReferencedElement with a parent field, so if this is non-nil the cast is correct
    if existingParent then
      return existingParent as structelem.StructureElement
    end
    local pending_parent: PendingParentElement = setmetatable({base = base}, pending_parent_metatable)
    return pending_parent
  end,
  exist = exist,
  exist_num = function(num: integer): boolean
    return exist('struct_num', num as string)
  end,
}
return labels