local write_pdf = require'luatagging-objectwriter'
local write_pdf_utils = require'luatagging-objectwriter-utils'
local names = write_pdf.names
local parenttree = require'luatagging-parenttree'
local rolemaps = require'luatagging-rolemaps'
local timing = require'luatagging-timing'

-- role_maps contains role maps for namespaces and elements which might not be used and therefore explicitly uses strings
-- and not objects.
local role_maps: {string | boolean: {string: string | {string}}} = rolemaps.role_maps

-- ref is only popuated if it has been written.
-- owner and ns should not be populated at the same time
local record Attributes
  ref: write_pdf.Ref
  owner: write_pdf.Name
  ns: string
  data: write_pdf.Dict
end
local record AttributeClass
  used_by_name: boolean
  attributes: {Attributes}
end

local attribute_registry: {write_pdf.Name: AttributeClass} = {}

local record Namespace
  ref: write_pdf.Ref
  ns: string
  types: {string: StructType}
  role_maps: {string: string | {string}}
end

local record StructType
  ns: Namespace
  name: write_pdf.Name
  role_map: StructType
  full_role_map: StructType
end

local namespaces: {string | boolean: Namespace}

local function resolve_rolemap(ns: Namespace, name: string): StructType
  local maps = ns.role_maps
  if not maps then
    return nil
  end
  local mapped = maps[name]
  if not mapped then
    return nil
  end
  local mapped_tag, mapped_ns: string, string | boolean
  if mapped is string then
    mapped_tag, mapped_ns = mapped, false
  else
    mapped_tag, mapped_ns = mapped[1], mapped[2]
  end
  return namespaces[mapped_ns].types[mapped_tag]
end

namespaces = setmetatable({}, {__index = function(t: {string | boolean: Namespace}, ns: string | boolean): Namespace
  local namespace: Namespace = {
    ns = (ns as string) or nil,
    role_maps = role_maps[ns],
  }
  if ns then
    namespace.ref = write_pdf_utils.delayed(function(): write_pdf.PdfObject
      local role_map = write_pdf.dict{}
      for tag, struct_type in pairs(namespace.types) do
        local mapped = struct_type.role_map
        if mapped then
          if mapped.ns.ref then
            role_map[tag] = write_pdf.array{mapped.name, mapped.ns.ref}
          else
            role_map[tag] = mapped.name
          end
        end
      end
      return write_pdf.dict {
        NS = ns as string,
        RoleMapNS = next(role_map) and role_map or nil,
      }
    end)
  end
  namespace.types = setmetatable({}, {__index = function(types: {string: StructType}, name: string): StructType
    local result: StructType = {
      ns = namespace,
      name = names[name],
      role_map = resolve_rolemap(namespace, name),
    }
    result.full_role_map = result.role_map and result.role_map.full_role_map or result.role_map or result
    types[name] = result
    return result
  end})
  t[ns] = namespace
  return namespace
end})

local function get_attribute_ref(self: Attributes): write_pdf.Ref
  if not self.ref then
    local data = self.data
    if self.ns then
      data.O = names.NSO
      data.NS = namespaces[self.ns].ref
    elseif self.owner then
      data.O = self.owner
    end
    self.ref = write_pdf_utils.indirect(data)
  end
  return self.ref
end

local interface ReferencedElement
  resolve: function(self): StructureElement
end
local interface StructureParent
  get_ref: function(self): write_pdf.Ref
  children: {StructureChild}
end
local interface StructureChild
  parent: StructureParent
  add_to_children_array: function(self, target: {write_pdf.PdfObject}, i: integer): integer
end

-- data is nil iff the StructureElement has been written out.
-- chilren is nil iff the structure elements children can no longer be changed.
-- children_ref is nil iff data is nil and children is not, indicating that
-- children are written separately.
local record StructureRoot is StructureChild, StructureParent
  children: {StructureChild}
end

global record StructureElement is StructureChild, StructureParent, ReferencedElement
  ref: write_pdf.Ref
  type: StructType
  data: StructureElementData
  children_ref: write_pdf.Ref
end

local record PendingElement is ReferencedElement
  resolved: StructureElement
end

local record StructureElementData
  id: string
  references: {ReferencedElement}
  page: write_pdf.Ref
  title: string
  language: string
  alt: string
  expanded: string
  actual_text: string
  associated_files: {write_pdf.Ref}
  attributes: {Attributes}
  classes: {write_pdf.Name}
  phonetic: Phonetic
end

local enum PhoneticAlphabet
end
-- TODO
local record Phonetic
  alphabet: PhoneticAlphabet
  phoneme: string
end

local record ObjectReference is StructureChild
  obj: write_pdf.Ref
  page: write_pdf.Ref
end

local structure_element_metatable: metatable<StructureElement> = {
  __index = StructureElement,
}
function StructureElement.add_to_children_array (structelem: StructureElement, children: {write_pdf.PdfObject}, i: integer): integer
  children[i] = structelem:get_ref()
  return i + 1
end
function StructureElement:resolve(): StructureElement
  return self
end

local object_reference_metatable: metatable<ObjectReference> = {
  __index = ObjectReference,
}
function ObjectReference.add_to_children_array (self: ObjectReference, children: {write_pdf.PdfObject}, i: integer): integer
  children[i] = write_pdf.dict{
    Type = names.OBJR,
    Obj = self.obj,
    Pg = self.page,
  }
  return i + 1
end

local structure_root_ref, write_structure_root = write_pdf_utils.late_ref()
local structure_root: StructureRoot = {
  get_ref = function(_: StructureRoot): write_pdf.Ref return structure_root_ref end,
  children = {},
}
local structure_stack: {StructureParent} = {}
local current_structure: StructureParent = structure_root
local unwritten_structures: {StructureElement} = {}

local function add_child(parent: StructureParent, child: StructureChild, first_child: boolean)
  assert(child.parent == nil)
  -- TODO: Handle already shipped
  -- TODO: Handle containment rules
  local children = parent.children
  if first_child then
    table.insert(children, 1, child)
  else
    children[#children + 1] = child
  end
  child.parent = parent
end

local function build_children(children: {StructureChild}): write_pdf.PdfObject
  local result: {write_pdf.PdfObject} = {}
  local i = 1
  for j = 1, #children do
    i = children[j]:add_to_children_array(result, i)
  end

  local length = #result
  if length == 0 then
    return nil
  elseif length == 1 then
    return result[1]
  else
    return write_pdf.array(result)
  end
end

local function get_ref_for_structelem(self: StructureElement): write_pdf.Ref
  local objnum = pdf.reserveobj()
  local ref = write_pdf.ref(objnum)
  function self:get_ref(): write_pdf.Ref
    return ref
  end
  self.ref = ref
  return ref
end

local function create_object_reference(obj: write_pdf.Ref, page: write_pdf.Ref, parent: StructureParent): ObjectReference
  local element: ObjectReference = setmetatable({
    obj = obj,
    page = page,
  }, object_reference_metatable)
  add_child(parent, element, false)
  return element
end

--[[
Known TODO:
 * Add IDs in IDTree
 * Attributes without classes
 * Handle Associated Files
]]

-- parent = false -> stashed
-- parent = true | parent = nil -> child of current element
local function create_structelem(tag: string, namespace: string, parent: boolean | StructureParent, first_child: boolean): StructureElement
  local structtype = namespaces[namespace or false].types[tag]
  local element: StructureElement = setmetatable({
    type = structtype,
    children = {},
    data = {},
    get_ref = get_ref_for_structelem,
  }, structure_element_metatable)
  unwritten_structures[#unwritten_structures + 1] = element
  if parent ~= false then
    add_child(parent == true and current_structure or (parent as StructureParent) or current_structure, element, first_child)
  end
  structure_stack[#structure_stack + 1] = current_structure
  current_structure = element
  return element
end

local function write_structelem(element: StructureElement)
  local struct_type = element.type
  local ns = struct_type.ns
  local data = element.data
  -- Unhandled fields in StructureElementData
  --   references: {ReferencedElement}
  --   associated_files: {any}
  --   attributes: {Attributes}
  -- end
  local classes: write_pdf.PdfObject
  if data.classes then
    if #data.classes == 1 then
      classes = data.classes[1]
      local attribute_dict = attribute_registry[classes]
      if attribute_dict then
        attribute_dict.used_by_name = true
      else
        tex.error(string.format("Unknown attribute class %s", classes[1]))
      end
    elseif #data.classes > 1 then
      classes = write_pdf.array(data.classes)
      for _, class in ipairs(classes) do
        local attribute_dict = attribute_registry[class]
        if attribute_dict then
          attribute_dict.used_by_name = true
        else
          tex.error(string.format("Unknown attribute class %s", class[1]))
        end
      end
    end
  end
  local references: write_pdf.Array
  if data.references then
    references = write_pdf.array{}
    for i, reference in ipairs(data.references) do
      local resolved = reference:resolve()
      if resolved then
        references[i] = resolved:get_ref()
      end
    end
  end
  local phonetic = data.phonetic
  local attributes: write_pdf.Array
  if data.attributes then
    attributes = write_pdf.array{}
    for i, attr in ipairs(data.attributes) do
      attributes[i] = get_attribute_ref(attr)
    end
  end
  pdf.immediateobj(element:get_ref()[1], write_pdf.to_pdf_string(write_pdf.dict{
    S = struct_type.name,
    P = element.parent and element.parent:get_ref(),
    ID = data.id,
    K = build_children(element.children), -- Needs to come before Pg to later allow filling this lazyly
    Ref = references and #references > 0 and references,
    Pg = data.page,
    NS = ns and ns.ref or nil,
    Lang = data.language,
    ActualText = data.actual_text,
    T = data.title,
    Alt = data.alt,
    E = data.expanded,
    PhoneticAlphabet = phonetic and phonetic.alphabet,
    Phoneme = phonetic and phonetic.phoneme,
    C = classes,
    AF = data.associated_files and write_pdf.array(data.associated_files),
    A = attributes and (#attributes > 1 and attributes or attributes[1]),
  }))
end

local function pop_structelem()
  assert(current_structure)
  local stack_depth = #structure_stack
  structure_stack[stack_depth], current_structure = nil, structure_stack[stack_depth]
end

luatexbase.add_to_callback('finish_pdffile', timing.write_structelem.wrap(function()
  for _, struct in ipairs(unwritten_structures) do
    write_structelem(struct)
  end
  local namespace_array, i = {}, 1
  for ns, namespace in pairs(namespaces) do
    if ns then
      namespace_array[i], i = namespace.ref, i + 1
    end
  end
  table.sort(namespace_array, function(a: write_pdf.Ref, b: write_pdf.Ref): boolean return a[1] < b[1] end)
  local classes = write_pdf.dict{}
  for classname, class in pairs(attribute_registry) do
    if class.used_by_name then
      local refs = write_pdf.array{}
      for j, attr in ipairs(class.attributes) do
        refs[j] = get_attribute_ref(attr)
      end
      classes[classname[1]] = refs
    end
  end
  if not next(classes) then
    classes = nil
  end
  table.sort(namespace_array, function(a: write_pdf.Ref, b: write_pdf.Ref): boolean return a[1] < b[1] end)
  write_structure_root(write_pdf.dict{
    Type = write_pdf.names.StructTreeRoot,
    K = build_children(structure_root.children),
    -- IDTree = TODO()
    ParentTree = parenttree.write_root(),
    -- RoleMap = TODO()
    ClassMap = classes,
    Namespaces = namespace_array[1] and write_pdf.array(namespace_array) or nil,
  })

  pdf.setcatalog((pdf.getcatalog() or '') .. write_pdf.to_pdf_dict_content{
    StructTreeRoot = structure_root:get_ref(),
    MarkInfo = write_pdf.dict {
      Marked = true,
    },
  })
end), 'luatagging.write_structelem')
luatexbase.declare_callback_rule('finish_pdffile', 'luatagging.write_structelem', 'before', 'luatagging.timing')

local record structelem
  type Attributes = Attributes
  type StructureChild = StructureChild
  type StructureParent = StructureParent
  type PendingElement = PendingElement
  type ReferencedElement = ReferencedElement
  type StructureElement = StructureElement
  type ObjectReference = ObjectReference

  create_structelem: function(tag: string, namespace: string, parent: boolean | StructureParent, first_child: boolean): StructureElement
  create_object_reference: function(obj: write_pdf.Ref, page: write_pdf.Ref, parent: StructureParent): ObjectReference
  pop_structelem: function()
  add_child: function(StructureParent, StructureChild, first_child?: boolean)
  current: function(): StructureParent
  register_attribute: function(write_pdf.Name, {Attributes})
  get_attribute: function(write_pdf.Name): {Attributes}
  _get_structure_stack: function(): {StructureParent}
  _used_namespaces: {string: Namespace}
end
local structelem_obj: structelem = {
  create_structelem = create_structelem,
  create_object_reference = create_object_reference,
  pop_structelem = pop_structelem,
  add_child = add_child,
  current = function(): StructureParent
    return current_structure
  end,
  register_attribute = function(name: write_pdf.Name, attrs: {Attributes})
    local class: AttributeClass = {
      attributes = attrs,
    }
    attribute_registry[name] = class
  end,
  get_attribute = function(name: write_pdf.Name): {Attributes}
    local class = attribute_registry[name]
    return class and class.attributes
  end,
  _get_structure_stack = function(): {StructureParent} return structure_stack end,
  _used_namespaces = namespaces,
}
return structelem_obj