local node = node
local traverse = node.traverse

local write_pdf = require'luatagging-objectwriter'
local structelem = require'luatagging-structelem'
local parenttree = require'luatagging-parenttree'
local timing = require'luatagging-timing'

local hlist_id = node.id'hlist'
local kern_id = node.id'kern'
local vlist_id = node.id'vlist'
local glue_id = node.id'glue'
local local_par_id = node.id'local_par'
local whatsit_id = node.id'whatsit'
local glyph_id = node.id'glyph'
local disc_id = node.id'disc'
local math_id = node.id'math'
local rule_id = node.id'rule'

local pdf_literal_sub = node.subtype'pdf_literal'

local mc_container_attribute = luatexbase.new_attribute'luatagging.mcc'
local type MarkedContentContainerId = integer
local record MarkedContent
  mcc: MarkedContentContainer
  page: integer
  mcid: integer
end
local record MarkedContentContainer is structelem.StructureChild
  mccid: MarkedContentContainerId
  tag: write_pdf.Name
  info: write_pdf.Dict
  mcs: {MarkedContent}
  parent: structelem.StructureElement
end
local marked_content_container_count: MarkedContentContainerId = 0
local marked_content_containers: {MarkedContentContainerId: MarkedContentContainer} = {}

function MarkedContentContainer.activate(self: MarkedContentContainer): MarkedContentContainer
  local old = tex.getattribute(mc_container_attribute)
  tex.setattribute('global', mc_container_attribute, self.mccid)
  return marked_content_containers[old]
end

local page_current_mcc: MarkedContentContainerId
local page_last_mcid = 0
local last_marked: node.Node
local page_mcids: {MarkedContent} = {}

local function add_mc_to_children_array(self: MarkedContentContainer, target: {write_pdf.PdfObject}, i: integer): integer
  for j, mc in ipairs(self.mcs) do
    target[i + j - 1] = write_pdf.dict{
      Type = write_pdf.names.MCR,
      Pg = write_pdf.ref(pdf.getpageref(mc.page)),
      MCID = mc.mcid,
    }
  end
  return i + #self.mcs
end

local function insert_mc_end_after(head: node.Node, n: node.Node): node.Node, node.Node
  local switch_whatsit = node.new(whatsit_id, pdf_literal_sub)
  switch_whatsit.attr = n.attr
  switch_whatsit.mode, switch_whatsit.data = 1, 'EMC'
  return node.insert_after(head, n, switch_whatsit)
end
local function insert_mc_begin_before(head: node.Node, n: node.Node, mccid: MarkedContentContainerId): node.Node, node.Node
  local mcc = assert(marked_content_containers[mccid])
  local switch_whatsit = node.new(whatsit_id, pdf_literal_sub)
  switch_whatsit.attr = n.attr
  local info = mcc.info
  local tag = mcc.tag
  local parent = mcc.parent
  if parent ~= nil then
    info = info or write_pdf.dict{}
    local mc: MarkedContent = {
      mcc = mcc,
      page = status.total_pages + 1,
      mcid = page_last_mcid,
    }
    mcc.mcs[#mcc.mcs + 1] = mc
    info.MCID = page_last_mcid
    page_last_mcid = page_last_mcid + 1
    page_mcids[page_last_mcid] = mc
    if not tag then
      tag = parent.type.name
    end
  end
  switch_whatsit.mode = 1
  if info then
    switch_whatsit.data = string.format('%s%sBDC', write_pdf.to_pdf_string(tag or write_pdf.names.Unknown), write_pdf.to_pdf_string(info))
  else
    switch_whatsit.data = string.format('%s BMC', write_pdf.to_pdf_string(tag or write_pdf.names.Unknown))
  end
  return node.insert_before(head, n, switch_whatsit)
end

local function process_marked(head: node.Node, marked_node: node.Node): node.Node
  local mcc: MarkedContentContainerId = node.get_attribute(marked_node, mc_container_attribute)
  if mcc == page_current_mcc then
    last_marked = marked_node
    return head
  end
  if last_marked ~= nil then
    head = insert_mc_end_after(head, last_marked)
  end
  page_current_mcc = mcc
  if not mcc then
    last_marked = nil
    return head
  end -- FIXME: Consider special mcc values, especially artifacts
  last_marked = marked_node
  head = insert_mc_begin_before(head, marked_node, mcc)
  return head
end

-- Lookup for link objnums which already got StructParents elsewhere.
local processed_annot_objects: {integer: boolean} = {}

local function process_start_link(head: node.Node, marked_node: node.Node): node.Node
  local mcc_id: MarkedContentContainerId = node.get_attribute(marked_node, mc_container_attribute)
  local mcc = mcc_id and marked_content_containers[mcc_id]
  if not mcc then return head end
  local objnum = marked_node.objnum
  if processed_annot_objects[objnum] then return head end
  local parent = mcc.parent
  if not mcc then return head end
  local type = parent.type
  type = type and type.rolemap or type
  processed_annot_objects[objnum] = true
  structelem.create_object_reference(write_pdf.ref(objnum), write_pdf.ref(pdf.getpageref(status.total_pages + 1)), parent)
  marked_node.link_attr = marked_node.link_attr .. write_pdf.to_pdf_dict_content { StructParent = parenttree.add(parent:get_ref()) }
  return head
end

-- How to process whatsit subtypes:
-- - false: Ignored for MC purposes
-- - function: Special handling
-- - nil: Unexpected whatsit
local whatsit_handling = {
  [node.subtype'open'] = false,
  [node.subtype'write'] = false,
  [node.subtype'close'] = false,
  [node.subtype'save_pos'] = false,
  [node.subtype'user_defined'] = false,
  [node.subtype'late_lua'] = process_marked,
  [node.subtype'special'] = process_marked,
  [node.subtype'pdf_literal'] = process_marked,
  [node.subtype'pdf_late_literal'] = process_marked,
  [node.subtype'pdf_refobj'] = false,
  [node.subtype'pdf_annot'] = false, -- process_annot,
  [node.subtype'pdf_start_link'] = process_start_link,
  [node.subtype'pdf_end_link'] = false,
  [node.subtype'pdf_dest'] = false, -- TODO: Consider automatic creation of structure destination
  -- [node.subtype'pdf_action'] = false, -- Should never be encountered.
  [node.subtype'pdf_thread'] = false,
  [node.subtype'pdf_start_thread'] = false,
  [node.subtype'pdf_colorstack'] = false,
  [node.subtype'pdf_setmatrix'] = false,
  [node.subtype'pdf_save'] = false,
  [node.subtype'pdf_restore'] = false,
  [node.subtype'pdf_link_state'] = false,
}

local process_hbox: function(node.HList)
local function process_vbox(box: node.VList)
  local original_head = box.head
  local head = original_head
  for n, id, sub in traverse(head) do
    if id == hlist_id then
      process_hbox(n as node.HList)
    elseif id == vlist_id then
      process_vbox(n as node.VList)
    elseif id == rule_id and sub ~= 3 then
      head = process_marked(head, n)
    elseif id == glue_id then
      local leader = n.leader
      if leader then
        -- TODO: Think about leaders with nested structure. Tricky, but also doesn't make much sense in the first place. I expect these mostly to be artifacts.
        head = process_marked(head, n)
      end
    elseif id == whatsit_id then
      local handler = whatsit_handling[sub]
      if handler then
        head = (handler as function(node.Node, node.Node): node.Node)(head, n)
      elseif handler == nil then
        texio.write_nl(string.format(
          "luatagging: Unsupported whatsit ignored in output: Subtype %i (%s)",
          sub,
          node.whatsits()[sub]
        ))
      end
    end
  end
  if head ~= original_head then
    box.head = head
  end
end

-- There is a limited number of nodes which can trigger output:
--  * rule (all subtypes except `none` and 0 dimensions)
--  * glyph (except virtual fonts which we ignore)
--  * disc (if rep not empty) (not really, but effectively we can't insert whatsit in the inner list so treat it as an atomic block)
--    * We could remove discs instead.
--  * glue if leader is set ... actually needs recursion into inner list, but repeating inner structure is problematic to treating as atomic is easier.
--  * whatsit (well, anything goes here. We could differentiate same subtypes)
local function process_hlist(head: node.Node): node.Node
  for n, id, sub in traverse(head) do
    if id == hlist_id then
      process_hbox(n as node.HList)
    elseif id == vlist_id then
      process_vbox(n as node.VList)
    elseif id == glue_id then
      local leader = n.leader
      if leader then
        -- TODO: Think about leaders with nested structure. Tricky, but also doesn't make much sense in the first place. I expect these mostly to be artifacts.
        head = process_marked(head, n)
      end
    elseif id == glyph_id or (id == rule_id and sub ~= 3) then
      head = process_marked(head, n)
    elseif id == disc_id then
      local disc = n as node.Disc
      if disc.replace ~= nil then
        head = process_marked(head, disc)
      end
    elseif id == whatsit_id then
      local handler = whatsit_handling[sub]
      if handler then
        head = (handler as function(node.Node, node.Node): node.Node)(head, n)
      elseif handler == nil then
        texio.write_nl(string.format(
          "luatagging: Unsupported whatsit ignored in output: Subtype %i (%s)",
          sub,
          node.whatsits()[sub]
        ))
      end
    end
  end
  return head
end

function process_hbox(box: node.HList)
  local original_head = box.head
  local head = process_hlist(original_head)
  if head ~= original_head then
    box.head = head
  end
end

-- local time_mc = 0
luatexbase.add_to_callback('pre_shipout_filter', timing.shipout_insert_mc.wrap(function(n: node.Node): boolean
  page_current_mcc = nil
  page_last_mcid = 0
  last_marked = nil
  page_mcids = {}

  if n.id == vlist_id then -- Very likely
    process_vbox(n as node.VList)
  elseif n.id == hlist_id then
    process_hbox(n as node.HList)
  else
    tex.error"Unexpected box type in pre_shipout_filter. This is a bug."
  end

  local parents = write_pdf.array{}
  for i, mc in ipairs(page_mcids) do
    parents[i] = mc.mcc.parent:get_ref()
  end
  ltx.__pdf.backend_ThisPage_gput(status.total_pages + 1, 'StructParents', parenttree.add(parents))
  return true
end), 'luatagging.insert_mc')
luatexbase.declare_callback_rule('pre_shipout_filter', 'luatagging.insert_spaces', 'before', 'luatagging.insert_mc') -- Needed to ensure that spaces get inserted before MCs get added.

local function new_mc_container(tag: write_pdf.Name, info: write_pdf.Dict, parent: StructureElement): MarkedContentContainer
  marked_content_container_count = marked_content_container_count + 1
  local container: MarkedContentContainer = setmetatable({
    mccid = marked_content_container_count,
    tag = tag,
    info = info,
    add_to_children_array = add_mc_to_children_array,
    mcs = {},
  }, {__index = MarkedContentContainer})
  if parent then
    structelem.add_child(parent, container)
  end
  marked_content_containers[marked_content_container_count] = container
  return container
end

function MarkedContentContainer.clone(self: MarkedContentContainer): MarkedContentContainer
  -- Think about what to copy.
  if not (self.parent or self.info) then
    return self
  end
  marked_content_container_count = marked_content_container_count + 1
  local container: MarkedContentContainer = setmetatable({
    mccid = marked_content_container_count,
    tag = self.tag,
    -- info = self.info,
    add_to_children_array = add_mc_to_children_array,
    mcs = {},
  }, {__index = MarkedContentContainer})
  if self.parent then
    structelem.add_child(self.parent, container)
  end
  marked_content_containers[marked_content_container_count] = container
  return container
end

local function new_mc_for_struct(struct: StructureElement, extra: write_pdf.Dict): MarkedContentContainer
  return new_mc_container(nil, extra or write_pdf.dict{}, struct)
end

local function new_mc_for_artifact(): MarkedContentContainer
  return new_mc_container(write_pdf.names.Artifact, nil, nil)
end

local record MarkedContentInterface
  type MarkedContentContainer = MarkedContentContainer

  new: function(struct: StructureElement, extra: write_pdf.Dict): MarkedContentContainer
  artifact: MarkedContentContainer
  attribute: integer
  processed_annot_objects: {integer: boolean}
end
local exported: MarkedContentInterface = {
  new = new_mc_for_struct,
  artifact = new_mc_for_artifact(),
  attribute = mc_container_attribute,
  processed_annot_objects = processed_annot_objects,
}
return exported