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'
local keyval = require'luatagging-lua3keyval'

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

local exp_after = token.create'exp_after:wN'
local expand_once_after_relax = {exp_after, token.create'scan_stop:', exp_after}

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

local function assign_context(field: string, scanner: function(boolean, string): any): function(boolean, any, string)
  return function(has_arg: boolean, context: any, key: string)
    local value = scanner(has_arg, key)
    if value == nil then return end
    context[field] = value
  end
end
local function append_clist(field: string, scanner: function(string): any, expand: boolean): function(boolean, any, string)
  local function appender(context: any)
    local list = context[field]
    if not list then
      list = {}
      context[field] = list
    end
    list[#list + 1] = scanner(field)
  end
  return function(has_arg: boolean, context: any, key: string)
    if not has_arg then
      return tex.error("Value required for " .. key)
    end
    keyval.clist(appender, context, expand)
  end
end
local function append_clist_flattened(field: string, scanner: function(string): {any}, expand: boolean): function(boolean, any, string)
  local function appender(context: any)
    local list = context[field]
    if not list then
      list = {}
      context[field] = list
    end
    local scanned = scanner(field)
    if scanned ==  nil then return end
    table.move(scanned, 1, #scanned, #list + 1, list)
  end
  return function(has_arg: boolean, context: any, key: string)
    if not has_arg then
      return tex.error("Value required for " .. key)
    end
    keyval.clist(appender, context, expand)
  end
end
local function append_one(field: string, scanner: function(string): any): function(boolean, any, string)
  return function(has_arg: boolean, context: any, key: string)
    if not has_arg then
      return tex.error("Value required for " .. key)
    end
    local list = context[field]
    if not list then
      list = {}
      context[field] = list
    end
    list[#list + 1] = scanner(field)
  end
end
local function with_default<T>(default: T, scanner: function(string): T): function(boolean, string): T
  return function(has_arg: boolean, key: string)
    if not has_arg then
      return default
    end
    return scanner(key)
  end
end
local function needs_value<T>(scanner: function(string): T): function(boolean, string): T
  return function(has_arg: boolean, key: string): T
    if not has_arg then
      return tex.error("Value required for " .. key)
    end
    return scanner(key)
  end
end
local function subcontext(field, scanner): function(boolean, any, string)
  return function(has_arg: boolean, context: any, key: string)
    local subcontext = context[field]
    if not subcontext then
      subcontext = {}
      context[field] = subcontext
    end
    scanner(has_arg, subcontext, key)
  end
end
local function scan_boolean(has_arg: boolean, key: string): boolean
  if not has_arg then return true end
  local arg = token.scan_argument()
  if arg == 'true' then
    return true
  elseif arg == 'false' then
    return false
  else
    tex.error("Unknown boolean value")
  end
end
local function scan_structnum(key: string): string
  token.scan_token() -- Scanning `{`
  local struct = labels.get_by_num(token.scan_int())
  token.scan_token() -- Scanning `}` -- TODO: Verify
  return struct
end
local function scan_string_e(key: string): string
  return token.scan_argument()
end
local function scan_string_n(key: string): string
  return token.scan_argument(false)
end
local function escaped(scanner: function(string): string): function(string): string
  return function(key: string): string
    local scanned = scanner(key)
    if not scanned then return end
    return escape_text(scanned)
  end
end
local function drop_empty(scanner: function(string): string): function(string): string
  return function(key: string): string
    local scanned = scanner(key)
    return scanned ~= '' and scanned or nil
  end
end
local function expand_once<T>(scanner: function(string): T): function(string): T
  return function(key: string): T
    token.put_next(expand_once_after_relax)
    token.scan_token()
    return scanner(key)
  end
end
local scan_string_o = expand_once(scan_string_n)

local function unsupported(has_arg: boolean, _: any, key: string)
  if has_arg then
    token.scan_argument(false)
  end
  tex.error("Unsupported structure element property " .. key)
end

local record StructKeyvalContext
  parent: structelem.StructureParent | boolean
  namespace: string
  tag: string
  data: structelem.StructureElementData
  first_child: boolean
  label: string
end

local tag_split_pattern = lpeg.C((1-lpeg.P'/')^0) * ('/' * lpeg.C((1-lpeg.P'/')^1))^-1 

local scan_text_nonempty_n = needs_value(escaped(drop_empty(scan_string_n)))
local scan_text_nonempty_o = needs_value(escaped(drop_empty(scan_string_o)))

local struct_keyval_handler = {
  label = assign_context('label', needs_value(drop_empty(scan_string_e))),
  firstkid = assign_context('first_child', scan_boolean),
  stash = function(has_arg: boolean, context: StructKeyvalContext, key: string)
    context.parent = (not scan_boolean(has_arg, key)) and nil
  end,
  parent = assign_context('parent', needs_value(scan_structnum)),
  ['parent-tag'] = unsupported,
  tag = function(has_arg: boolean, context: StructKeyvalContext, key: string)
    local combined_tag = scan_string_e(key)
    local tag, namespace = tag_split_pattern:match(combined_tag)
    context.tag = tag
    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
      context.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
      context.namespace = namespace
    end
  end,
  title = subcontext('data', assign_context('title', scan_text_nonempty_n)),
  ['title-o'] = subcontext('data', assign_context('title', scan_text_nonempty_o)),
  alt = subcontext('data', assign_context('alt', scan_text_nonempty_o)),
  alttext = subcontext('data', assign_context('alt', scan_text_nonempty_o)),
  actualtext = subcontext('data', assign_context('actual_text', scan_text_nonempty_o)),
  phoneme = subcontext('data', subcontext('phonetic', assign_context('phoneme', scan_text_nonempty_o))),
  lang = subcontext('data', assign_context('language', needs_value(drop_empty(scan_string_e)))),
  ref = subcontext('data', append_clist('references', function(key: string): structelem.ReferencedElement
    local label = scan_string_n(key)
    return labels.get('label', label)
  end)),
  E = subcontext('data', assign_context('expanded', scan_text_nonempty_o)),
  AF = unsupported,
  AFref = subcontext('data', append_clist('associated_files', function(key: string): write_pdf.Ref
    local content = scan_string_e(key)
    return write_pdf.raw(content) as write_pdf.Ref
  end)),
  AFinline = unsupported,
  ['AFinline-o'] = unsupported,
  texsource = subcontext('data', append_one('associated_files', function(key: string): write_pdf.Ref
    local content = scan_string_o(key)
    return associated_files.register(content, write_pdf.names["application/x-tex"], write_pdf.names.Supplement, "source.tex", nil, nil)
  end)),
  mathml = unsupported,
  ['attribute-class'] = subcontext('data', append_clist('classes', function(key: string): write_pdf.Name
    return names[drop_empty(scan_string_n)(key)]
  end, true)),
  attribute = subcontext('data', append_clist_flattened('attributes', function(key: string): structelem.Attributes
    local classname = names[drop_empty(scan_string_n)(key)]
    if classname == nil then return end
    local class_attributes = structelem.get_attribute(classname)
    if not class_attributes then
      tex.error(string.format("Unknown attribute %s", classname[1]))
      return
    end
    return class_attributes
  end, true)),
}

timed_tex_function('StructBeginKV', function()
  local context: StructKeyvalContext = {}
  keyval.keyval(struct_keyval_handler, context)
  local struct = structelem.create_structelem(context.tag, context.namespace, context.parent, context.first_child)
  if context.data then
    local data = struct.data
    for k, v in pairs(context.data) do
      data[k] = v
    end
  end
end, 'protected')

local record McKeyvalContext
  artifact: boolean
  tag: write_pdf.Name
end
local mc_keyval_handler = {
  artifact = assign_context('artifact', scan_boolean),
  tag = assign_context('tag', scan_string_e),
}

local artifact_name = names.Artifact
-- ActivateMc / ActivateMcArtifact / SetMcProperty hybrid with keyval.
-- This should be almost exactly \tag_mc_begin:n
timed_tex_function('BeginMc', function()
  local context: McKeyvalContext = {}
  keyval.keyval(mc_keyval_handler, context)
  local mcc
  if context.artifact then
    if context.tag and context.tag ~= 'Artifact' then
      tex.error"Changing the tag of an artifact MC segment is not allowed."
    end
    marked_content.artifact:activate()
  else
    local current = structelem.current()
    if current is StructureElement then
      mcc = marked_content.new(current, nil)
    else
      return tex.error"Can not start MC segment outside any structure"
    end
    if context.tag then
      mcc.tag = names[context.tag]
    end
    mcc:activate()
  end
end, 'protected')