module:depends("muc");

local jid_resource = require "util.jid".resource;
local st = require "util.stanza";

local prefixes = module:get_option_set("muc_inject_mentions_prefixes", {})
local suffixes = module:get_option_set("muc_inject_mentions_suffixes", {})
local enabled_rooms = module:get_option("muc_inject_mentions_enabled_rooms", nil)
local disabled_rooms = module:get_option("muc_inject_mentions_disabled_rooms", nil)
local mention_delimiters = module:get_option_set("muc_inject_mentions_mention_delimiters",  {" ", "", "\n", "\t"})
local append_mentions = module:get_option("muc_inject_mentions_append_mentions", false)
local strip_out_prefixes = module:get_option("muc_inject_mentions_strip_out_prefixes", false)
local reserved_nicks = module:get_option("muc_inject_mentions_reserved_nicks", false)
local use_bare_jid = module:get_option("muc_inject_mentions_use_bare_jid", true)
local prefix_mandatory = module:get_option("muc_inject_mentions_prefix_mandatory", false)
local reserved_nicknames = {}

local reference_xmlns = "urn:xmpp:reference:0"

local function update_reserved_nicknames(event)
    local room, data, jid = event.room.jid, event.data, event.jid
    load_room_reserved_nicknames(event.room)
    local nickname = (data or {})["reserved_nickname"]

    if nickname then
        reserved_nicknames[room][nickname] = jid
    else
        local nickname_to_remove
        for _nickname, _jid in pairs(reserved_nicknames[room]) do
            if _jid == jid then
                nickname_to_remove = _nickname
                break
            end
        end
        if nickname_to_remove then
            reserved_nicknames[room][nickname_to_remove] = nil
        end
    end
end

function load_room_reserved_nicknames(room)
    if not reserved_nicknames[room.jid] then
        reserved_nicknames[room.jid] = {}
        for jid, data in pairs(room._affiliations_data or {}) do
            local reserved_nickname = data["reserved_nickname"]
            if reserved_nicknames then
                reserved_nicknames[room.jid][reserved_nickname] = jid
            end
        end
    end
end

local function get_jid(room, nickname)
    local bare_jid = reserved_nicknames[room.jid][nickname]
    if bare_jid and use_bare_jid then
        return bare_jid
    end

    if bare_jid and not use_bare_jid then
        return room.jid .. "/" .. nickname
    end
end

local function get_participants(room)
    if not reserved_nicks then
        local occupants = room._occupants
        local key, occupant = next(occupants)
        return function ()
            while occupant do -- luacheck: ignore
                local nick = jid_resource(occupant.nick);
                local bare_jid = occupant.bare_jid
                key, occupant = next(occupants, key)
                return bare_jid, nick
            end
        end
    else
        local generator = room:each_affiliation()
        local jid, _, affiliation_data = generator(nil, nil)
        return function ()
           while jid do
                local bare_jid, nick = jid, (affiliation_data or {})["reserved_nickname"]
                jid, _, affiliation_data = generator(nil, bare_jid)
                if nick then
                    return bare_jid, nick
                end
           end
        end
    end
end

local function add_mention(mentions, bare_jid, first, last, prefix_indices, has_prefix)
    if strip_out_prefixes then
        if has_prefix then
            table.insert(prefix_indices, first-1)
        end
        first = first - #prefix_indices
        last = last - #prefix_indices
    end
    mentions[first] = {bare_jid=bare_jid, first=first, last=last}
end

local function get_client_mentions(stanza)
    local has_mentions = false
    local client_mentions = {}

    for element in stanza:childtags("reference", reference_xmlns) do
        if element.attr.type == "mention" then
            local key = tonumber(element.attr.begin) + 1 -- count starts at 0
            client_mentions[key] = {bare_jid=element.attr.uri, first=element.attr.begin, last=element.attr["end"]}
            has_mentions = true
        end
    end

    return has_mentions, client_mentions
end

local function is_room_eligible(jid)
    if not enabled_rooms and not disabled_rooms then return true; end

    if enabled_rooms then
        for _, _jid in ipairs(enabled_rooms) do
            if _jid == jid then
                return true
            end
        end
        return false
    end

    if disabled_rooms then
        for _, _jid in ipairs(disabled_rooms) do
            if _jid == jid then
                return false
            end
        end
        return true
    end
end

local function has_nick_prefix(body, first)
    -- There are no configured prefixes
    if not prefixes or #prefixes < 1 then return false end

    -- Prefix must have a space before it,
    -- be the first character of the body
    -- or be the first character after a new line
    if not mention_delimiters:contains(body:sub(first - 2, first - 2)) then
        return false
    end

    local prefix = body:sub(first - 1, first - 1)
    for _, _prefix in ipairs(prefixes) do
        if prefix == _prefix then
            return true
        end
    end

    return false
end

local function has_nick_suffix(body, last)
    -- There are no configured suffixes
    if not suffixes or #suffixes < 1 then return false end

    -- Suffix must have a space after it,
    -- be the last character of the body
    -- or be the last character before a new line
    if not mention_delimiters:contains(body:sub(last + 2, last + 2)) then
        return false
    end

    local suffix = body:sub(last+1, last+1)
    for _, _suffix in ipairs(suffixes) do
        if suffix == _suffix then
            return true
        end
    end

    return false
end

local function search_mentions(room, body, client_mentions)
    load_room_reserved_nicknames(room)
    local mentions, prefix_indices = {}, {}
    local current_word = ""
    local current_word_start
    for i = 1, #body+1 do
        local char = body:sub(i,i)
        -- Mention delimiter found, current_word is completed now
        if mention_delimiters:contains(char) and current_word_start then
            -- Check for nickname without prefix
            local jid = get_jid(room, current_word)
            if jid then
                if not prefix_mandatory then
                    add_mention(mentions, jid, current_word_start, i - 1, prefix_indices, false)
                end
            else
                -- Check for nickname with affixes
                local prefix = prefixes:contains(current_word:sub(1,1))
                local suffix = suffixes:contains(current_word:sub(-1))
                if prefix and suffix then
                    jid = get_jid(room, current_word:sub(2, -2))
                    if jid then
                        add_mention(mentions, jid, current_word_start + 1, i - 2, prefix_indices, true)
                    end
                elseif prefix then
                    jid = get_jid(room, current_word:sub(2))
                    if jid then
                        add_mention(mentions, jid, current_word_start + 1, i - 1, prefix_indices, true)
                    end
                elseif suffix and not prefix_mandatory then
                    jid = get_jid(room, current_word:sub(1, -2))
                    if jid then
                        add_mention(mentions, jid, current_word_start, i - 2, prefix_indices, false)
                    end
                end
            end

            current_word = ""
            current_word_start = nil
        elseif not mention_delimiters:contains(char) then
            current_word_start = current_word_start or i
            current_word = current_word .. char
        end
    end

    return mentions, prefix_indices
end

local function muc_inject_mentions(event)
    local room, stanza = event.room, event.stanza;
    local body = stanza:get_child_text("body")

    if not body or #body < 1 then return; end

    -- Inject mentions only if the room is configured for them
    if not is_room_eligible(room.jid) then return; end

    -- Only act on messages that do not include mentions
    -- unless configuration states otherwise.
    local has_mentions, client_mentions = get_client_mentions(stanza)
    if has_mentions and not append_mentions then return; end

    local mentions, prefix_indices = search_mentions(room, body, client_mentions)
    for _, mention in pairs(mentions) do
        -- https://xmpp.org/extensions/xep-0372.html#usecase_mention
        stanza:tag(
            "reference", {
                xmlns=reference_xmlns,
                begin=tostring(mention.first - 1), -- count starts at 0
                ["end"]=tostring(mention.last),
                type="mention",
                uri="xmpp:" .. mention.bare_jid,
            }
        ):up()
    end

    if strip_out_prefixes then
        local body_without_prefixes = ""
        local from = 0
        if #prefix_indices > 0 then
            for _, prefix_index in ipairs(prefix_indices) do
                body_without_prefixes = body_without_prefixes .. body:sub(from, prefix_index-1)
                from = prefix_index + 1
            end
            body_without_prefixes = body_without_prefixes .. body:sub(from, #body)

            -- Replace original body containing prefixes
            stanza:maptags(
                function(tag)
                    if tag.name ~= "body" then
                        return tag
                    end
                    return st.stanza("body"):text(body_without_prefixes)
                end
            )
        end
    end
end

module:hook("muc-occupant-groupchat", muc_inject_mentions)
module:hook("muc-set-affiliation", update_reserved_nicknames)