"""
transpiler/runtime_snippets.py — Collection of Luau helper functions for injection.

Each snippet is a standalone Luau string. 
Dependencies are tracked so that if one helper is used, its dependencies are also injected.
"""

# Mapping of helper name -> (source_code, dependencies)
SNIPPETS: dict[str, tuple[str, list[str]]] = {
    "py_bool": (
        """local function py_bool(value)
    if value == nil or value == false or value == 0 or value == 0.0 or value == "" then return false end
    if type(value) == "table" and next(value) == nil then return false end
    return true
end""",
        []
    ),

    "py_len": (
        """local function py_len(x)
    local t = type(x)
    if t == "string" then
        return string.len(x)
    elseif t == "table" then
        if x._is_set then return x._count end
        return #x
    end
    error("py_len: unsupported type '" .. t .. "'", 2)
end""",
        []
    ),

    "py_str": (
        """local function py_str(x)
    if x == nil then return "None" end
    if x == true then return "True" end
    if x == false then return "False" end
    return tostring(x)
end""",
        []
    ),

    "py_int": (
        """local function py_int(x)
    local n = tonumber(x)
    if n == nil then error("py_int: invalid literal: '" .. tostring(x) .. "'", 2) end
    return math.modf(n)
end""",
        []
    ),

    "py_float": (
        """local function py_float(x)
    local n = tonumber(x)
    if n == nil then error("py_float: could not convert: '" .. tostring(x) .. "'", 2) end
    return n
end""",
        []
    ),

    "py_range": (
        """local function py_range(a, b, c)
    local start, stop, step
    if b == nil then start, stop, step = 0, a, 1
    elseif c == nil then start, stop, step = a, b, 1
    else start, stop, step = a, b, c end
    local result = {}
    local current = start
    if step > 0 then
        while current < stop do
            table.insert(result, current)
            current = current + step
        end
    elseif step < 0 then
        while current > stop do
            table.insert(result, current)
            current = current + step
        end
    end
    return result
end""",
        []
    ),

    "py_print": (
        """local function py_print(...)
    local args = {...}
    local parts = {}
    for _, v in ipairs(args) do
        table.insert(parts, py_str(v))
    end
    print(table.concat(parts, "\\t"))
end""",
        ["py_str"]
    ),

    "py_contains": (
        """local function py_contains(collection, value)
    local t = type(collection)
    if t == "string" then
        return string.find(collection, value, 1, true) ~= nil
    elseif t == "table" then
        if collection._is_set then return collection.items[value] == true end
        for _, v in ipairs(collection) do if v == value then return true end end
        if collection[value] ~= nil then return true end
    end
    return false
end""",
        []
    ),

    "py_index": (
        """local function py_index(tbl, idx)
    if idx < 0 then return tbl[#tbl + 1 + idx] end
    return tbl[idx + 1]
end""",
        []
    ),

    "py_append": (
        """local function py_append(list, value)
    table.insert(list, value)
end""",
        []
    ),

    "py_pop": (
        """local function py_pop(list, idx)
    if idx == nil then return table.remove(list) end
    if idx < 0 then idx = #list + 1 + idx else idx = idx + 1 end
    return table.remove(list, idx)
end""",
        []
    ),

    "py_slice": (
        """local function py_slice(tbl, start, stop, step)
    local n = #tbl
    step = step or 1
    if start == nil then start = step > 0 and 0 or (n - 1)
    elseif start < 0 then start = math.max(0, n + start) end
    if stop == nil then stop = step > 0 and n or -1
    elseif stop < 0 then stop = math.max(0, n + stop) end
    local result = {}
    local i = start
    if step > 0 then
        while i < stop do
            table.insert(result, tbl[i + 1])
            i = i + step
        end
    elseif step < 0 then
        while i > stop do
            table.insert(result, tbl[i + 1])
            i = i + step
        end
    end
    return result
end""",
        []
    ),

    "py_iter": (
        """local function py_iter(iterable)
    local t = type(iterable)
    if t == "table" then
        if iterable._is_set then
            local f, s, var = pairs(iterable.items)
            return function(_, prev)
                local k = f(s, prev)
                if k ~= nil then return 0, k end
                return nil
            end, s, var
        end
        return ipairs(iterable)
    end
    return function() end
end""",
        []
    ),

    "py_enumerate": (
        """local _ENUM_POOL = {}
local function py_enumerate(list, start)
    start = start or 0
    local result = {}
    for i, v in ipairs(list) do
        local entry = table.remove(_ENUM_POOL) or {0, nil}
        entry[1] = start + (i - 1)
        entry[2] = v
        table.insert(result, entry)
    end
    -- Note: We can't easily return items to pool here without knowing when user is done.
    -- But for AAA games, we provide a py_release(result) helper.
    return result
end""",
        []
    ),

    "py_release": (
        """local function py_release(list)
    for _, entry in ipairs(list) do
        if type(entry) == "table" and #entry == 2 then -- Check if it's a 2-element array
            entry[1] = nil -- Clear references
            entry[2] = nil
            table.insert(_ENUM_POOL, entry)
        end
    end
end""",
        []
    ),
    
    # Dict helpers
    "py_keys": (
        """local function py_keys(dict)
    local res = {}
    for k in pairs(dict) do table.insert(res, k) end
    return res
end""",
        []
    ),
    
    "py_values": (
        """local function py_values(dict)
    local res = {}
    for _, v in pairs(dict) do table.insert(res, v) end
    return res
end""",
        []
    ),

    "py_items": (
        """local function py_items(dict)
    local res = {}
    for k, v in pairs(dict) do
        local entry = table.remove(_ENUM_POOL) or {nil, nil}
        entry[1] = k
        entry[2] = v
        table.insert(res, entry)
    end
    return res
end""",
        []
    ),
    # String helpers
    "py_split": (
        """local function py_split(str, sep)
    if sep == nil then
        local res = {}
        for w in string.gmatch(str, "%S+") do table.insert(res, w) end
        return res
    end
    local res = {}
    local pattern = "(.-)" .. string.gsub(sep, "([^%w])", "%%%1")
    local last_end = 1
    for part, pos in string.gmatch(str, pattern .. "()") do
        table.insert(res, part)
        last_end = pos
    end
    table.insert(res, string.sub(str, last_end))
    return res
end""",
        []
    ),

    # Set helpers
    "py_set_new": (
        """local function py_set_new(iterable)
    local s = { _is_set = true, items = {}, _count = 0 }
    if iterable then
        for _, v in py_iter(iterable) do
            if not s.items[v] then
                s.items[v] = true
                s._count = s._count + 1
            end
        end
    end
    return s
end""",
        ["py_iter"]
    ),

    "py_set_add": (
        """local function py_set_add(s, item)
    if not s.items[item] then
        s.items[item] = true
        s._count = s._count + 1
    end
end""",
        []
    ),

    "py_set_discard": (
        """local function py_set_discard(s, item)
    if s.items[item] then
        s.items[item] = nil
        s._count = s._count - 1
    end
end""",
        []
    ),

    "py_set_clear": (
        """local function py_set_clear(s)
    s.items = {}
    s._count = 0
end""",
        []
    ),

    "py_set_union": (
        """local function py_set_union(s1, s2)
    local s = py_set_new()
    for k in pairs(s1.items) do py_set_add(s, k) end
    for k in pairs(s2.items) do py_set_add(s, k) end
    return s
end""",
        ["py_set_new", "py_set_add"]
    ),

    "py_set_intersection": (
        """local function py_set_intersection(s1, s2)
    local s = py_set_new()
    for k in pairs(s1.items) do
        if s2.items[k] then py_set_add(s, k) end
    end
    return s
end""",
        ["py_set_new", "py_set_add"]
    ),

    "py_set_difference": (
        """local function py_set_difference(s1, s2)
    local s = py_set_new()
    for k in pairs(s1.items) do
        if not s2.items[k] then py_set_add(s, k) end
    end
    return s
end""",
        ["py_set_new", "py_set_add"]
    ),

    "py_eq": (
        """local function py_eq(a, b)
    if a == b then return true end
    local ta = type(a)
    if ta ~= type(b) then return false end
    if ta ~= "table" then return false end
    
    if a._is_set or b._is_set then
        if not (a._is_set and b._is_set) then return false end
        if a._count ~= b._count then return false end
        for k in pairs(a.items) do
            if not b.items[k] then return false end
        end
        return true
    end

    -- Deep table comparison
    local count_a, count_b = 0, 0
    for k, v in pairs(a) do
        count_a = count_a + 1
        if not py_eq(v, b[k]) then return false end
    end
    for _ in pairs(b) do count_b = count_b + 1 end
    return count_a == count_b
end""",
        []
    ),

    # re module shims
    "py_re_match": (
        """local function py_re_match(pattern, str)
    local s, e = string.find(str, "^" .. pattern)
    if s then
        return { group = function(self, n) return string.sub(str, s, e) end } -- Simplistic Match object
    end
    return nil
end""",
        []
    ),

    "py_re_search": (
        """local function py_re_search(pattern, str)
    local s, e = string.find(str, pattern)
    if s then
        return { group = function(self, n) return string.sub(str, s, e) end }
    end
    return nil
end""",
        []
    ),

    "py_re_sub": (
        """local function py_re_sub(pattern, repl, str)
    return (string.gsub(str, pattern, repl))
end""",
        []
    ),

    # json module shims
    "py_json_dumps": (
        """local function py_json_dumps(obj)
    local HttpService = game:GetService("HttpService")
    return HttpService:JSONEncode(obj)
end""",
        []
    ),

    "py_json_loads": (
        """local function py_json_loads(str)
    local HttpService = game:GetService("HttpService")
    return HttpService:JSONDecode(str)
end""",
        []
    ),

    # math shims
    "py_math_radians": (
        """local function py_math_radians(x)
    return math.rad(x)
end""",
        []
    ),
    "py_math_degrees": (
        """local function py_math_degrees(x)
    return math.deg(x)
end""",
        []
    ),

    # random shims
    "py_random_random": (
        """local function py_random_random()
    return math.random()
end""",
        []
    ),
    "py_random_randint": (
        """local function py_random_randint(a, b)
    return math.random(a, b)
end""",
        []
    ),
    "py_pcall": (
        """local function py_pcall(func, ...)
    return pcall(func, ...)
end""",
        []
    ),

    # Buffer helpers (Phase Q: AAA Scale)
    "py_buffer_create": (
        """local function py_buffer_create(size)
    return buffer.create(size)
end""",
        []
    ),

    "py_buffer_readi8": (
        """local function py_buffer_readi8(b, offset)
    return buffer.readi8(b, offset)
end""",
        []
    ),

    "py_buffer_writei8": (
        """local function py_buffer_writei8(b, offset, value)
    buffer.writei8(b, offset, value)
end""",
        []
    ),

    "py_buffer_readf32": (
        """local function py_buffer_readf32(b, offset)
    return buffer.readf32(b, offset)
end""",
        []
    ),

    "py_buffer_writef32": (
        """local function py_buffer_writef32(b, offset, value)
    buffer.writef32(b, offset, value)
end""",
        []
    ),
    # collections module shims (Phase Q2: AAA Stdlib)
    "py_collections_deque_new": (
        """local function py_collections_deque_new(iterable)
    local d = { _is_deque = true, head = 1, tail = 0, items = {} }
    if iterable then
        for _, v in py_iter(iterable) do
            d.tail = d.tail + 1
            d.items[d.tail] = v
        end
    end
    return d
end""",
        ["py_iter"]
    ),
    "py_collections_deque_append": (
        """local function py_collections_deque_append(d, v)
    d.tail = d.tail + 1
    d.items[d.tail] = v
end""",
        []
    ),
    "py_collections_deque_appendleft": (
        """local function py_collections_deque_appendleft(d, v)
    d.head = d.head - 1
    d.items[d.head] = v
end""",
        []
    ),
    "py_collections_deque_pop": (
        """local function py_collections_deque_pop(d)
    if d.head > d.tail then error("pop from an empty deque", 2) end
    local v = d.items[d.tail]
    d.items[d.tail] = nil
    d.tail = d.tail - 1
    return v
end""",
        []
    ),
    "py_collections_deque_popleft": (
        """local function py_collections_deque_popleft(d)
    if d.head > d.tail then error("pop from an empty deque", 2) end
    local v = d.items[d.head]
    d.items[d.head] = nil
    d.head = d.head + 1
    return v
end""",
        []
    ),
    "py_collections_counter_new": (
        """local function py_collections_counter_new(iterable)
    local c = { _is_counter = true, items = {} }
    setmetatable(c, {
        __index = function(self, key) return self.items[key] or 0 end,
        __newindex = function(self, key, val) self.items[key] = val end
    })
    if iterable then
        for _, v in py_iter(iterable) do
            c.items[v] = (c.items[v] or 0) + 1
        end
    end
    return c
end""",
        ["py_iter"]
    ),
    "py_collections_defaultdict_new": (
        """local function py_collections_defaultdict_new(default_factory)
    local d = {}
    setmetatable(d, {
        __index = function(self, key)
            local val = default_factory()
            rawset(self, key, val)
            return val
        end
    })
    return d
end""",
        []
    ),
    "py_collections_namedtuple_new": (
        """local function py_collections_namedtuple_new(typename, field_names)
    local fields = {}
    if type(field_names) == "string" then
        for w in string.gmatch(field_names, "%S+") do table.insert(fields, w) end
    else
        fields = field_names
    end
    return function(...)
        local args = {...}
        local inst = {}
        for i, name in ipairs(fields) do
            inst[name] = args[i]
            inst[i] = args[i]
        end
        return inst
    end
end""",
        []
    ),
    # heapq module shims (Phase Q2: AAA Stdlib)
    "py_heapq_heappush": (
        """local function py_heapq_heappush(heap, item)
    table.insert(heap, item)
    local child = #heap
    while child > 1 do
        local parent = math.floor(child / 2)
        if heap[child] < heap[parent] then
            heap[child], heap[parent] = heap[parent], heap[child]
            child = parent
        else
            break
        end
    end
end""",
        []
    ),
    "py_heapq_heappop": (
        """local function py_heapq_heappop(heap)
    if #heap == 0 then error("index out of range", 2) end
    local last = table.remove(heap)
    if #heap == 0 then return last end
    local root = heap[1]
    heap[1] = last
    local parent = 1
    while true do
        local child = 2 * parent
        if child > #heap then break end
        if child + 1 <= #heap and heap[child + 1] < heap[child] then
            child = child + 1
        end
        if heap[child] < heap[parent] then
            heap[child], heap[parent] = heap[parent], heap[child]
            parent = child
        else
            break
        end
    end
    return root
end""",
        []
    ),

    # bisect module shims (Phase Q2: AAA Stdlib)
    "py_bisect_bisect_left": (
        """local function py_bisect_bisect_left(a, x, lo, hi)
    lo = lo or 0
    hi = hi or #a
    while lo < hi do
        local mid = math.floor((lo + hi) / 2)
        if a[mid + 1] < x then lo = mid + 1 else hi = mid end
    end
    return lo
end""",
        []
    ),
    "py_bisect_bisect_right": (
        """local function py_bisect_bisect_right(a, x, lo, hi)
    lo = lo or 0
    hi = hi or #a
    while lo < hi do
        local mid = math.floor((lo + hi) / 2)
        if x < a[mid + 1] then hi = mid else lo = mid + 1 end
    end
    return lo
end""",
        []
    ),
    "py_mult": (
        """local function py_mult(a, b)
    if type(a) == "number" and type(b) == "number" then
        return a * b
    end

    if type(a) == "table" and type(b) == "number" then
        local res = {}
        for i = 1, b do
            for _, v in ipairs(a) do
                table.insert(res, v)
            end
            for k, v in pairs(a) do
                if type(k) ~= "number" then
                    res[k] = v
                end
            end
        end
        return res
    end

    if type(b) == "table" and type(a) == "number" then
        return py_mult(b, a)
    end

    -- Safe fallback for Roblox types (Vector3, CFrame, etc)
    return a * b
end""",
        []
    ),
}

def get_used_snippets(used_names: set[str]) -> str:
    """
    Builds the final Luau header by recursively resolving dependencies
    of all used helper functions.
    """
    to_inject = set()
    stack = list(used_names)
    
    while stack:
        name = stack.pop()
        if name in SNIPPETS and name not in to_inject:
            to_inject.add(name)
            _, deps = SNIPPETS[name]
            stack.extend(deps)
            
    if not to_inject:
        return ""
        
    names = sorted(list(to_inject))
    lines = [f"local {', '.join(names)}"]
    
    for name in names:
        code, _ = SNIPPETS[name]
        # remove "local " from the start of the standard 'local function name()' snippet
        if code.startswith("local "):
            code = code[6:]
        lines.append(code)
        
    return "\n".join(lines) + "\n"
