-- 有序映射类，支持自定义比较器
-- 内部使用哈希表 (_data) 存储键值对，数组 (_keys) 维护有序的键列表
local OrderedMap = {}
OrderedMap.__index = OrderedMap

-- 创建新映射
-- @param comparator 比较器，可以是：
--   - 字符串 "asc" 或 "<" 表示升序（默认）
--   - 字符串 "desc" 或 ">" 表示降序
--   - 函数 cmp(a, b) 返回 true 表示 a 应该排在 b 前面
function OrderedMap.new(comparator)
    -- 创建元表指向自身，使所有实例共享方法
    local self = setmetatable({}, OrderedMap)
    self._data = {}          -- 键 -> 值的哈希表
    self._keys = {}          -- 有序的键列表
    self._cmp = nil          -- 比较函数

    -- 根据参数初始化比较器
    if comparator == nil or comparator == "asc" or comparator == "<" then
        self._cmp = function(a, b) return a < b end   -- 默认升序
    elseif comparator == "desc" or comparator == ">" then
        self._cmp = function(a, b) return a > b end   -- 降序
    elseif type(comparator) == "function" then
        self._cmp = comparator                         -- 自定义函数
    else
        error("comparator must be a function, 'asc', 'desc', '<', or '>'")
    end
    return self
end

-- 内部方法：二分查找键应插入的位置（保持顺序）
-- 返回插入位置的索引（第一个使得 cmp(key, _keys[i]) 为假的索引）
function OrderedMap:_findInsertPos(key)
    local keys = self._keys      -- 缓存局部变量，提高访问速度
    local cmp = self._cmp
    local low, high = 1, #keys
    while low <= high do
        local mid = math.floor((low + high) / 2)
        -- 若 key 应排在 keys[mid] 前面，则继续在前半部分查找
        if cmp(key, keys[mid]) then
            high = mid - 1
        else
            low = mid + 1        -- 否则在后半部分查找（包括相等的情况）
        end
    end
    return low                   -- 返回应插入的位置
end

-- 插入或更新键值对
-- 注意：不允许存储 nil 值，若传入 nil 则视为删除该键
function OrderedMap:set(key, value)
    -- 处理值为 nil 的情况：调用 remove 删除键并返回
    if value == nil then
        self:remove(key)
        return
    end

    local data = self._data      -- 缓存局部变量
    local keys = self._keys

    -- 如果键已存在，仅更新值，不改变顺序
    if data[key] ~= nil then
        data[key] = value
        return
    end

    -- 插入新键：先存储值，再找到合适位置插入键
    data[key] = value
    local pos = self:_findInsertPos(key)
    table.insert(keys, pos, key)  -- 在 pos 位置插入键，后续元素自动后移
end

-- 获取键对应的值（若键不存在返回 nil）
function OrderedMap:get(key)
    return self._data[key]
end

-- 删除键值对
function OrderedMap:remove(key)
    local data = self._data
    -- 若键不存在，直接返回
    if data[key] == nil then return end
    data[key] = nil               -- 从哈希表中删除值

    local keys = self._keys
    -- 遍历有序键列表，找到该键并删除
    for i, k in ipairs(keys) do
        if k == key then
            table.remove(keys, i) -- 删除后，后续元素自动前移
            return
        end
    end
end

-- 返回有序键的数组（副本，避免外部修改）
function OrderedMap:keys()
    local copy = {}
    for i, k in ipairs(self._keys) do
        copy[i] = k
    end
    return copy
end

-- 返回迭代器函数，用于 for 循环按顺序遍历所有键值对
-- 用法：for key, value in map:ipairs() do ... end
function OrderedMap:ipairs()
    local keys = self._keys        -- 缓存局部变量
    local data = self._data
    local i = 0
    -- 闭包迭代器，每次调用返回下一个键值对
    return function()
        i = i + 1
        local key = keys[i]
        if key then
            return key, data[key]
        end
    end
end

-- 返回键值对数量
function OrderedMap:size()
    return #self._keys
end

-- 清空所有数据
function OrderedMap:clear()
    self._data = {}
    self._keys = {}
end

-- 支持 Lua 5.2+ 的 __pairs 元方法，使 pairs(map) 能按顺序迭代
-- 直接复用 ipairs 方法（ipairs 本身返回迭代器）
OrderedMap.__pairs = OrderedMap.ipairs

-- 为兼容性提供别名，建议显式调用 ipairs
OrderedMap.pairs = OrderedMap.ipairs

-- 示例使用（仅当直接运行此文件时执行）
if ... then
    -- 升序（默认）
    local map = OrderedMap.new()  -- 或 "asc"
    map:set("b", 2)
    map:set("a", 1)
    map:set("c", 3)
    for k, v in map:ipairs() do
        print(k, v)  -- 输出 a 1, b 2, c 3
    end

    -- 降序
    local map2 = OrderedMap.new("desc")
    map2:set(5, "five")
    map2:set(1, "one")
    map2:set(3, "three")
    for k, v in map2:ipairs() do
        print(k, v)  -- 输出 5 five, 3 three, 1 one
    end

    -- 自定义比较器：按字符串长度排序
    local map3 = OrderedMap.new(function(a, b) return #a < #b end)
    map3:set("apple", "fruit")
    map3:set("banana", "fruit")
    map3:set("pear", "fruit")
    for k, v in map3:ipairs() do
        print(k, v)  -- 输出 pear fruit, apple fruit, banana fruit (长度升序)
    end
end

return OrderedMap