local write_pdf = require'luatagging-objectwriter'
local structelem = require'luatagging-structelem'
local marked_content = require'luatagging-marked-content'
local rolemaps = require'luatagging-rolemaps'
local pdf_escape = require'luatagging-pdf-escape'
local escape_text = pdf_escape.to_pdf_encoding
local labels = require'luatagging-labels'
local associated_files = require'luatagging-associated-files'
local parenttree = require'luatagging-parenttree'
local timing = require'luatagging-timing'
require'luatagging-charprocessing'
require'luatagging-kv'

local expl3_return_bool = {
  [true] = token.create'prg_return_true:',
  [false] = token.create'prg_return_false:',
}

local names = write_pdf.names
local namespace_lookup = rolemaps.namespace_lookup

local function timed_tex_function(macro_name: string, handler: function(), ...)
  local func = luatexbase.new_luafunction(macro_name)
  token.set_lua(macro_name, func, ...)
  lua.get_functions_table()[func] = timing[macro_name].wrap(handler)
end

timed_tex_function('StructBegin', function()
  local first: boolean
  local parent: boolean | structelem.StructureParent = not token.scan_keyword'stashed'
  if parent then
    first = token.scan_keyword'first'
    if token.scan_keyword'parent' then
      local kind = token.scan_string()
      if kind == 'struct_num' then
        parent = labels.get_by_num(token.scan_int())
      end
    end
  end
  local namespace = token.scan_keyword'namespace' and token.scan_argument() or nil
  local tag = token.scan_argument()
  if namespace then
    local found_namespace = namespace_lookup[namespace]
    if not found_namespace then
      tex.error(string.format("Unknown namespace %q", namespace))
      found_namespace = namespace_lookup.user
    end
    namespace = found_namespace
  else
    namespace = rolemaps.default_namespaces[tag]
    if not namespace then
      tex.error(string.format("No default namespace for tag %q", tag))
      namespace = namespace_lookup.user
    end
  end
  local _struct = structelem.create_structelem(tag, namespace, parent, first)
end, 'protected')

local unescaped_slashed_name = '/' * ((1-lpeg.S'/,')^1 / function(s) return names[s] end)
local unescaped_slashed_name_list = lpeg.Ct((unescaped_slashed_name * (',' * unescaped_slashed_name)^0)^-1) * -1

local properties: {string: function(StructureElement)} = {
  Lang = function(struct: StructureElement)
    local language = token.scan_argument()
    assert(struct.data).language = language
  end,
  ActualText = function(struct: StructureElement)
    local text = token.scan_argument()
    assert(struct.data).actual_text = escape_text(text)
  end,
  Title = function(struct: StructureElement)
    local text = token.scan_argument()
    assert(struct.data).title = escape_text(text)
  end,
  Alt = function(struct: StructureElement)
    local text = token.scan_argument()
    assert(struct.data).alt = escape_text(text)
  end,
  Expanded = function(struct: StructureElement)
    local text = token.scan_argument()
    assert(struct.data).expanded = text
  end,
  Phoneme = function(struct: StructureElement)
    local text = token.scan_argument()
    local phonetic = assert(struct.data).phonetic
    if not phonetic then
      phonetic = {}
      struct.data.phonetic = phonetic
    end
    phonetic.phoneme = escape_text(text)
  end,
  Dest = function(struct: StructureElement)
    local label = token.scan_argument()
    labels.assign(struct, 'dest', label)
  end,
  Label = function(struct: StructureElement)
    local label = token.scan_argument()
    labels.assign(struct, 'label', label)
  end,
  Attribute = function(struct: StructureElement)
    local classname = names[token.scan_argument()]
    local data = assert(struct.data)
    local class_attributes = structelem.get_attribute(classname)
    if not class_attributes then
      tex.error(string.format("Unknown attribute %s", classname[1]))
      return
    end
    local attributes = data.attributes
    if not attributes then
      attributes = {}
      data.attributes = attributes
    end
    table.move(class_attributes, 1, #class_attributes, #attributes + 1, attributes)
  end,
  ["Attribute-inline"] = function(struct: StructureElement)
    local content = token.scan_argument()
    local data = assert(struct.data)
    local attributes = data.attributes
    if not attributes then
      attributes = {}
      data.attributes = attributes
    end
    attributes[#attributes + 1] = {data = write_pdf.dict{[write_pdf.raw] = content}}
  end,
  ["Attribute-class"] = function(struct: StructureElement)
    local class = names[token.scan_argument()]
    local classes = assert(struct.data).classes
    if not classes then
      classes = {}
      struct.data.classes = classes
    end
    classes[#classes + 1] = class
  end,
  ["Attribute-class-list"] = function(struct: StructureElement)
    local classes = assert(unescaped_slashed_name_list:match(token.scan_argument()))
    local data = assert(struct.data)
    struct.data.classes = classes
  end,
  ["Ref-num"] = function(struct: StructureElement)
    local num = assert(math.tointeger(token.scan_argument()))
    local references = assert(struct.data).references
    if not references then
      references = {}
      struct.data.references = references
    end
    references[#references + 1] = labels.get_by_num(num)
  end,
  ["Ref-label"] = function(struct: StructureElement)
    local label = token.scan_argument()
    local references = assert(struct.data).references
    if not references then
      references = {}
      struct.data.references = references
    end
    references[#references + 1] = labels.get('label', label)
  end,
  ["Ref-dest"] = function(struct: StructureElement)
    local label = token.scan_argument()
    local references = assert(struct.data).references
    if not references then
      references = {}
      struct.data.references = references
    end
    references[#references + 1] = labels.get('dest', label)
  end,
  ["Ref-dest-parent"] = function(struct: StructureElement)
    local label = token.scan_argument()
    local references = assert(struct.data).references
    if not references then
      references = {}
      struct.data.references = references
    end
    references[#references + 1] = labels.get_parent(labels.get('dest', label))
  end,
  ["AF-raw"] = function(struct: StructureElement)
    local content = token.scan_argument()
    local af = assert(struct.data).associated_files
    if not af then
      af = {}
      struct.data.associated_files = af
    end
    af[#af + 1] = write_pdf.raw(content) as write_pdf.Ref -- The type is a lie, but in practice the raw value corresponds to a ref
  end,
  ["AF-texsource"] = function(struct: StructureElement)
    local content = token.scan_argument()
    local af = assert(struct.data).associated_files
    if not af then
      af = {}
      struct.data.associated_files = af
    end
    af[#af + 1] = associated_files.register(content, write_pdf.names["application/x-tex"], write_pdf.names.Supplement, "source.tex", nil, nil)
  end,
  Tag = function(struct: StructureElement)
    local tag = token.scan_argument() -- TODO: Verify compatibility and maybe check containment rules.
    struct.type = assert(struct.type.ns.types[tag])
  end,
  [""] = function(struct: StructureElement)
    token.scan_argument()
  end,
}

local mc_properties: {string: function(marked_content.MarkedContentContainer)} = {
  Tag = function(mcc: marked_content.MarkedContentContainer)
    local tag = token.scan_argument()
    mcc.tag = write_pdf.names[tag]
  end,
}

timed_tex_function('SetStructProperty', function()
  local property = token.scan_string()
  local handler = properties[property]
  if handler == nil then
    tex.error(string.format("Unsupported structure property %q(%s) ignored.", property, type(property)))
    token.scan_argument()
    return
  end
  local struct = assert(structelem.current()) as StructureElement -- FXME: Handle root
  return handler(struct)
end, 'protected')

timed_tex_function('SetStructPropertyForStructNum', function()
  local struct = labels.get_by_num(token.scan_int())
  if not struct then
    token.scan_string()
    token.scan_argument()
    tex.error("Can't set properties on unknown structure number")
    return
  end
  local property = token.scan_string()
  local handler = properties[property]
  if handler == nil then
    token.scan_argument()
    tex.error(string.format("Unsupported structure property %q ignored.", property))
    return
  end
  return handler(struct)
end, 'protected')

timed_tex_function('StructEnd', function()
  structelem.pop_structelem()
end, 'protected')

timed_tex_function('SetMcProperty', function()
  local property = token.scan_string()
  local handler = mc_properties[property]
  if handler == nil then
    tex.error(string.format("Unsupported MC property %q(%s) ignored.", property, type(property)))
    token.scan_argument()
    return
  end
  local mcc = assert(structelem.current()) as marked_content.MarkedContentContainer -- FXME: Handle inactive or static artifact
  return handler(mcc)
end, 'protected')

timed_tex_function('ActivateMc', function()
  local current = structelem.current()
  if current is StructureElement then
    marked_content.new(current, nil):activate()
  else
    tex.error"Can not start MC segment outside any structure"
  end
end, 'protected')

local suspended_mc_stack = {}

timed_tex_function('PushMc', function()
  suspended_mc_stack[#suspended_mc_stack + 1] = marked_content.artifact:activate() or marked_content.artifact
end, 'protected')

timed_tex_function('PopMc', function()
  local popped = suspended_mc_stack[#suspended_mc_stack]
  if not popped then
    return tex.error"luatagging: Unable to pop MC without previous push"
  end
  suspended_mc_stack[#suspended_mc_stack] = nil
  -- TODO: Think about this.
  popped:clone():activate()
end, 'protected')

timed_tex_function('ActivateMcArtifact', function()
  marked_content.artifact:activate()
end, 'protected')

timed_tex_function('RegisterAttribute', function()
  local name = token.scan_argument()
  local content = token.scan_argument()
  structelem.register_attribute(names[name], {{data = write_pdf.dict{[write_pdf.raw] = content}}})
end, 'protected')

timed_tex_function('AddTagNS', function()
  local name = token.scan_argument()
  local ns = token.scan_argument()
  local mapped = token.scan_argument()
  local mapped_ns = token.scan_argument()
  local namespace = namespace_lookup[ns]
  if not namespace then
    tex.error(string.format("Unknown namespace name %s", ns))
    namespace = namespace_lookup.user
  end
  local mapped_namespace = namespace_lookup[mapped_ns]
  if not mapped_namespace then
    tex.error(string.format("Unknown namespace name %s", mapped_ns))
    namespace = namespace_lookup.pdf
  end
  rolemaps.default_namespaces[name] = namespace
  rolemaps.add_mapping(name, namespace, mapped, mapped_namespace)
  -- structelem.register_attribute(names[name], {{data = write_pdf.dict{[write_pdf.raw] = content}}})
end, 'protected')

timed_tex_function('GetCurrentStructTag', function()
  local struct = assert(structelem.current()) as StructureElement -- FIXME: Handle root
  tex.sprint(struct.type.name[1])
end)

timed_tex_function('SetUserNamespace', function()
  local ns = token.scan_argument()
  local old_ns = namespace_lookup.user
  namespace_lookup.user = ns
end, 'protected')

timed_tex_function('IfTagKnown', function()
  local tag = token.scan_argument()
  token.put_next(expl3_return_bool[rolemaps.default_namespaces[tag] ~= nil])
end, 'protected')

timed_tex_function('IfStructNumDefined', function()
  local struct_num = token.scan_int()
  token.put_next(expl3_return_bool[labels.exist_num(struct_num)])
end, 'protected')

timed_tex_function('GetStructNum', function()
  local struct = assert(structelem.current()) as StructureElement -- FIXME: Handle root
  tex.write(tostring(labels.assign_num(struct)))
end)

timed_tex_function('UseStructByNum', function()
  local child_num = token.scan_int()
  local child = labels.get_by_num(child_num)
  local parent = structelem.current()
  if child.parent then
    tex.error"Illegal attempt to attach structure element multiple times."
    return
  end
  structelem.add_child(parent, child)
end, 'protected')

-- Weird one:
--  Take the currently active structure element (which should be stashed and add it as a child
--  to the first sect or part ancestor of the previously active structure element.
timed_tex_function('AttachToMarginparParent', function()
  local child = structelem.current() as structelem.StructureElement -- If cast is wrong then the structure stack is empty and we report an error
  local structure_stack = structelem._get_structure_stack()
  local surrounding = structure_stack[#structure_stack]
  if not surrounding or child.parent then error"Misuse" end
  local namespace = rawget(structelem._used_namespaces, namespace_lookup.pdf2)
  local sect_type = namespace and rawget(namespace.types, 'Sect')
  local part_type = namespace and rawget(namespace.types, 'Part')
  local found: boolean
  while surrounding and (surrounding as structelem.StructureElement).type do
    local current_structure_element = surrounding as structelem.StructureElement
    local current_type = current_structure_element.type
    local role_mapped = current_type.role_map or current_type
    if role_mapped == sect_type or role_mapped == part_type then
      found = true
      break
    end
    surrounding = current_structure_element.parent
  end
  if not surrounding then
    tex.error"\\marginpar used in a stashed structure. Unable to identity parent. The structure tree will be wrong."
  end
  structelem.add_child(found and surrounding or structure_stack[2], child) -- StructureStack[2] is the current document structure. It is used as a fallback.
end, 'protected')

timed_tex_function('ResetStructureAttributeToCurrentValueForBox', function()
  local box_id = token.scan_int()
  local box = tex.box[box_id]
  if not box then return end
  local attr = marked_content.attribute
  local attr_value = tex.getattribute(attr)
  for n in node.traverse(box.list) do
    node.set_attribute(n, attr, attr_value)
  end
end, 'protected')

local char_given_cmd = assert(token.command_id'char_given')
timed_tex_function('GetCurrentStructureReference', function()
  token.put_next(token.create(assert(structelem.current()):get_ref()[1], char_given_cmd))
end)

timed_tex_function('GetStructClassClist', function()
  local child_num = token.scan_int()
  local child = labels.get_by_num(child_num)
  local classes = assert(child.data).classes
  if not classes then return end
  local string_classes = {}
  for i, class in ipairs(classes) do
    if i ~= 1 then string_classes[3*i-3] = ',' end
    string_classes[3*i-2] = '/'
    string_classes[3*i-1] = class[1] -- sic. We use the (unescaped) class[1] instead of the escaped class[2] since we need to read these back as unescaped values.
  end
  tex.sprint(string_classes)
end)

do
  local annot_parent: structelem.StructureParent
  local last_struct_parent: integer
  timed_tex_function('AssignAndGetStructParent', function()
    local struct = assert(structelem.current()) as structelem.StructureParent
    if annot_parent then
      if annot_parent ~= struct then
        tex.error"Multiple requests for the same StructParent with different structure parents. This is a bug."
      else
        -- I would like to warn here, but this is happening rather often and not really an issue.
        -- texio.write_nl"Multiple requests for the same StructParent."
      end
    else
      annot_parent = struct
      last_struct_parent = parenttree.add(struct:get_ref())
    end
    tex.write(tostring(last_struct_parent))
  end)
  timed_tex_function('RecordObjectForStructParent', function()
    local object_str = token.scan_argument()
    local object = string.match(object_str, '%s*(%d+)%s+0%s+R')
    if not object then
      return tex.error"Invalid first argument to \\tag_struct_insert_annot:nn"
    end
    local object_ref = write_pdf.ref(assert(math.tointeger(object)))
    local provided_struct_parent = token.scan_argument()
    local struct = assert(structelem.current()) as structelem.StructureParent
    if not annot_parent then
      return tex.error"\\tag_struct_insert_annot:nn used without \\tag_struct_parent_int:."
    end
    if struct ~= annot_parent then
      annot_parent = nil
      return tex.error"\\tag_struct_insert_annot:nn used in different structure element than \\tag_struct_parent_int:."
    end
    annot_parent = nil
    if tostring(last_struct_parent) ~= provided_struct_parent then
      annot_parent = nil
      return tex.error"Wrong second argument to \\tag_struct_insert_annot:nn."
    end
    marked_content.processed_annot_objects[object_ref[1]] = true
    local whatsit = node.new('whatsit', 'late_lua')
    whatsit.data = function()
      structelem.create_object_reference(object_ref, write_pdf.ref(pdf.getpageref(status.total_pages + 1)), struct)
    end
    node.write(whatsit)
  end, 'protected')
end
package.preload['luamml-structelemwriter'] = loadfile(kpse.find_file'luatagging-luamml-structelemwriter.lua')