local timing = require'luatagging-timing'
local marked_content = require'luatagging-marked-content'
local identifiers = fonts.hashes.identifiers
local node = node
local direct = node.direct
local getlist, setlist, todirect, traverse, traverse_list, traverse_id, traverse_char, traverse_glyph, end_of_math, node_new, insert_before, setfont, getfont, setwidth, getwidth, getdisc, getattributelist, setattributelist, set_attribute = direct.getlist, direct.setlist, direct.todirect, direct.traverse, direct.traverse_list, direct.traverse_id, direct.traverse_char, direct.traverse_glyph, direct.end_of_math, direct.new, direct.insert_before, direct.setfont, direct.getfont, direct.setwidth, direct.getwidth, direct.getdisc, direct.getattributelist, direct.setattributelist, direct.set_attribute
assert(traverse_id)
local properties = direct.get_properties_table()
local record CharProcessingProperties
  luatagging_spacefont: integer
  luatagging_is_soft_hyphen: boolean
end
local hlist_id = node.id'hlist'
local vlist_id = node.id'vlist'
local glyph_id = node.id'glyph'
local glue_id = node.id'glue'
local math_id = node.id'math'
local disc_id = node.id'disc'
local discardable = {
  [glue_id] = true,
  [node.id'penalty'] = true,
  [node.id'kern'] = true,
  -- [math_id] = true, -- Intentionally omitted since math nodes need special handling for our purposes.
}

local explicit_disc_sub = 1
local regular_disc_sub = 3

local type State = integer
local initial: State = 0
local after_glyph: State = 1
local after_glue: State = 2

-- The space code should be pretty much equivalent to the tagpdf code, this implementation is roughly 4 times faster in benchmarks.

-- We try to find non-zero glue nodes which have a glyph on at least one side.
-- Discardable nodes in between are ignored, but if a glue is found the first discardable node is marked insted of the actual glue.
-- TODO: Cleanup, but be careful about performance. This is very much on the hot path. 
-- Potentially write it as a tail call based state machine? Needs benchmarking.
local function mark_spaces(head: integer): boolean
  local iter, iter_state, n = traverse(head)
  while true do
    local id, sub
    n, id, sub = iter(iter_state, n)
    if not n then return true end

    if id == glyph_id then
      local glyph = n
      local first_discardable
      while true do
        n, id, sub = iter(iter_state, n)
        if not n then return true end
        if id == glue_id and getwidth(n) > 0 then
          first_discardable = first_discardable or n
          local props = properties[first_discardable]
          if not props then
            props = {}
            properties[first_discardable] = props
          end
          props.luatagging_spacefont = getfont(glyph)
          while true do
            n, id, sub = iter(iter_state, n)
            if not n then return true end
            if not discardable[id] then break end
          end
          break
        elseif discardable[id] then
          if not first_discardable then
            first_discardable = n
          end
        else
          break
        end
      end
      iter_state, n = n, nil
    elseif discardable[id] then
      local first_discardable, first_glue = n, nil
      if id == glue_id and getwidth(n) > 0 then
        first_glue = n
      else
        while true do
          n, id, sub = iter(iter_state, n)
          if not n then return true end
          if id == glue_id and getwidth(n) > 0 then
            first_glue = n
            break
          end
          if not discardable[id] then break end
        end
      end
      if first_glue then
        while true do
          n, id, sub = iter(iter_state, n)
          if not n then return true end
          if id == glyph_id then
            local props = properties[first_discardable]
            if not props then
              props = {}
              properties[first_discardable] = props
            end
            props.luatagging_spacefont = getfont(n)
            while true do
              n, id, sub = iter(iter_state, n)
              if not n then return true end
              if not discardable[id] then break end
            end
            break
          end
          if not discardable[id] then break end
        end
      end
      iter_state, n = n, nil
    elseif id == math_id then
      if sub == 0 then
        -- Start of math. Skip to end
        local stop_math = end_of_math(n)
        -- The stop_math is a valid break point if it's directly followed by glue so in that case it needs to be marked
        -- as space if we would consider the glue as space (since the glue would be discarded at a linebreak)
        if stop_math then
          n = stop_math
          n, id, sub = iter(iter_state, n)
          if not n then return true end
          if id == glue_id and getwidth(n) > 0 then
            local first_glue = n
            while true do
              n, id, sub = iter(iter_state, n)
              if not n then return true end
              if id == glyph_id then
                local props = properties[stop_math]
                if not props then
                  props = {}
                  properties[stop_math] = props
                end
                props.luatagging_spacefont = getfont(n)
                while true do
                  n, id, sub = iter(iter_state, n)
                  if not n then return true end
                  if not discardable[id] then break end
                end
                break
              end
              if not discardable[id] then break end
            end
          end
          iter_state, n = n, nil
        else
          texio.write_nl("luatagging: Failed to find end of math segment. Tagging of spaces might be wrong.")
        end
      else
        texio.write_nl("luatagging: Unexpected end of math segment. Tagging of spaces might be wrong.")
      end
    end
  end
end

luatexbase.add_to_callback('pre_shaping_filter', timing.mark_spaces.wrap(function(n: node.Node): node.Node | boolean
  mark_spaces(todirect(n))
  return true
end), 'luatagging.mark_spaces')

local font_spaceglyph = setmetatable({}, {__index = function(t, fid)
  local glyph = 32
  glyph = glyph or false
  t[fid] = glyph
  return glyph
end})

local insert_spaces_at_shipout_hbox: function(integer): integer
local function insert_spaces_at_shipout_vlist(head: integer)
  for n, id, _, list in traverse_list(head) do
    if id == hlist_id then
      local new_list = insert_spaces_at_shipout_hlist(list)
      if new_list ~= list then
        setlist(n, new_list)
      end
    elseif id == vlist_id then
      insert_spaces_at_shipout_vlist(list)
    else
      error'Invalid id in traverse_list'
    end
  end
end
function insert_spaces_at_shipout_hlist(head: integer): integer
  local iter, iter_state, n = traverse(head)
  while true do
    local id, sub
    n, id, sub = iter(iter_state, n)
    if not n then return head end

    if id == hlist_id then
      local list = getlist(n)
      local new_list = insert_spaces_at_shipout_hlist(list)
      if new_list ~= list then
        setlist(n, new_list)
      end
    elseif id == vlist_id then
      insert_spaces_at_shipout_vlist(getlist(n))
    else
      local props = properties[n]
      local spacefont = props and props.luatagging_spacefont
      if spacefont then
        local marked = n
        while true do
          if id == glue_id then break end
          n, id, sub = iter(iter_state, n)
          if not n then break end
        end
        local space_glyph = node_new(glyph_id)
        local gid = font_spaceglyph[spacefont]
        if gid then
          setfont(space_glyph, spacefont, gid)
        else
          error'TODO: Find good defaults'
          setfont(space_glyph, something, defaulty)
        end
        if n then
          -- We have a glue. Adjust it.
          setwidth(n, getwidth(n) - getwidth(space_glyph))
        else
          -- TODO: We don't have a glue. This should only happen at the end of a line. No need to adjust here I think (Maybe for r2l?)
          n = marked
        end
        setattributelist(space_glyph, getattributelist(n))
        head = insert_before(head, n, space_glyph)
      end
    end
  end
  return head
end

luatexbase.add_to_callback('pre_shipout_filter', timing.shipout_insert_spaces.wrap(function(n: node.Node): boolean
  insert_spaces_at_shipout_vlist(getlist(todirect(n))) -- Omit the getlist to make this a bit safer if called with very direct lists
  return true
end), 'luatagging.insert_spaces')

-- softhyphen processing: This is the same as in tagpdf, we just replicate it here to set the right attributes and use direct mode for good measure.

local hyphen_char = 0x2D
local soft_hyphen_char = 0xAD
local softhyphen_fonts = setmetatable({}, {__index = function(t, fid)
  local fdir = identifiers[fid]
  local format = fdir and fdir.format
  local result = (format == 'opentype' or format == 'truetype')
  local characters = fdir and fdir.characters
  result = result and (characters and characters[soft_hyphen_char]) ~= nil
  t[fid] = result
  return result
end})

-- A pre shaping callback to mark hyphens as being hyphenation hyphens.
-- This runs before shaping to avoid affecting hyphens moved into
-- discretionaries during shaping.
local function process_softhyphen_pre(head: integer): boolean
  -- if softhyphenbool.mode ~= truebool.mode then return true end
  for disc, sub in traverse_id(disc_id, head) do
    if sub == explicit_disc_sub or sub == regular_disc_sub then
      local pre = getdisc(disc)
      for n in traverse_char(pre) do
        local props = properties[n]
        if not props then
          props = {}
          properties[n] = props
        end
        props.luatagging_is_soft_hyphen = true
      end
    end
  end
  return true
end

-- Finally do the actual replacement after shaping. No checking for double processing here
-- since the operation is idempotent.
local mc_attribute = marked_content.attribute
local artifact_value = marked_content.artifact.mccid
local function process_softhyphen_post(head, _context, _dir): boolean
  -- if softhyphenbool.mode ~= truebool.mode then return true end
  for disc, sub in traverse_id(disc_id, head) do
    local pre = getdisc(disc)
    for n, ch, fid in traverse_glyph(pre) do
      if ch == hyphen_char then
        local props = properties[n]
        if props and props.luatagging_is_soft_hyphen then
          -- if nodegetattribute(n,softhyphenattribute) and softhyphen_fonts[fid] then
          --   n.char = soft_hyphen_char
          --   props.glyph_info = nil
          -- else
            set_attribute(n, mc_attribute, artifact_value)
          -- end
        end
      end
    end
  end
  return true
end

luatexbase.add_to_callback('pre_shaping_filter', timing.mark_softhyphen.wrap(function(n: node.Node): boolean
  return process_softhyphen_pre(todirect(n))
end), 'luatagging.rewrite-softhyphen')
luatexbase.add_to_callback('post_shaping_filter', timing.process_softhyphen.wrap(function(n: node.Node): boolean
  return process_softhyphen_post(todirect(n))
end), 'luatagging.rewrite-softhyphen')