#!/usr/bin/env lua
-- gen_atom_offsets.lua — hand-rolled, no patterns, no regex, no region pragmas.
--
-- Finds every `MipsAtom_(name) { ... }` declaration in the given sources,
-- counts the words in each body using the WORD_COUNT manifest, computes
-- branch offsets for atom_label(name) / atom_offset(tag, name) markers,
-- and writes one header per source into <source_dir>/gen/<basename>.offsets.h
--
-- Generated header layout (per source):
--   #pragma region <basename>
--   #undef atom_offset
--   #define atom_offset(tag, name)  atom_offset_##tag##_##name
--   // --- atom: <name> (<n> words) ---
--   #define atom_offset_<tag>_<target>  (N)        // preprocessor form
--   #undef atom_offset_<tag>_<target>              // (so enum can reuse)
--   enum {
--       atom_offset_<tag>_<target> = N,            // C enum form
--   };
--   #define atom_offset_<tag>_<target>  (N)        // re-define for preprocessor
--   #pragma endregion <basename>
--
-- Usage:
--   lua gen_atom_offsets.lua <metadata.h> <source1> [source2 ...]

-- ============================================================
-- Character classification
-- ============================================================

local function is_space(c)
    return c == " " or c == "\t" or c == "\n" or c == "\r" or c == "\v" or c == "\f"
end

local function is_alpha(c)
    if not c or #c == 0 then return false end
    if c >= "a" and c <= "z" then return true end
    if c >= "A" and c <= "Z" then return true end
    return c == "_"
end

local function is_digit(c)
    return c and c >= "0" and c <= "9"
end

local function is_alnum(c)
    return is_alpha(c) or is_digit(c)
end

-- ============================================================
-- I/O
-- ============================================================

local function read_file(path)
    local f = io.open(path, "r")
    if not f then error("Cannot open " .. path) end
    local content = f:read("*a")
    f:close()
    return content
end

local function write_file(path, content)
    local f = io.open(path, "w")
    if not f then error("Cannot write " .. path) end
    f:write(content)
    f:close()
end

local function ensure_dir(path)
    os.execute('mkdir -p "' .. path .. '"')
end

-- ============================================================
-- String primitives
-- ============================================================

local function trim(s)
    local a = 1
    while a <= #s and is_space(s:sub(a, a)) do a = a + 1 end
    local b = #s
    while b >= a and is_space(s:sub(b, b)) do b = b - 1 end
    return s:sub(a, b)
end

local function starts_with(s, prefix)
    if #s < #prefix then return false end
    for i = 1, #prefix do
        if s:sub(i, i) ~= prefix:sub(i, i) then return false end
    end
    return true
end

local function ends_with(s, suffix)
    if #s < #suffix then return false end
    local off = #s - #suffix
    for i = 1, #suffix do
        if s:sub(off + i, off + i) ~= suffix:sub(i, i) then return false end
    end
    return true
end

local function find_byte(haystack, target, start)
    for i = start or 1, #haystack do
        if haystack:sub(i, i) == target then return i end
    end
    return nil
end

local function dirname(path)
    local last_sep = 0
    for i = 1, #path do
        local c = path:sub(i, i)
        if c == "/" or c == "\\" then last_sep = i end
    end
    if last_sep == 0 then return "." end
    return path:sub(1, last_sep - 1)
end

local function basename_no_ext(path)
    local last_sep = 0
    for i = 1, #path do
        local c = path:sub(i, i)
        if c == "/" or c == "\\" then last_sep = i end
    end
    local a = last_sep + 1
    local last_dot = #path + 1
    for i = #path, a, -1 do
        if path:sub(i, i) == "." then last_dot = i; break end
    end
    return path:sub(a, last_dot - 1)
end

local function to_upper(s)
    local out = ""
    for i = 1, #s do
        local code = string.byte(s, i)
        if code >= 97 and code <= 122 then
            out = out .. string.char(code - 32)
        else
            out = out .. s:sub(i, i)
        end
    end
    return out
end

local function to_alnum_underscore(s)
    local out = ""
    for i = 1, #s do
        local c = s:sub(i, i)
        if is_alnum(c) then out = out .. c
        else out = out .. "_" end
    end
    return out
end

local function pad_right(s, width)
    while #s < width do s = s .. " " end
    return s
end

-- ============================================================
-- Lexer helpers
-- ============================================================

local function skip_ws_and_comments(source, i)
    local len = #source
    while i <= len do
        local c = source:sub(i, i)
        if is_space(c) then
            i = i + 1
        elseif c == "/" and source:sub(i+1, i+1) == "/" then
            while i <= len and source:sub(i, i) ~= "\n" do i = i + 1 end
        elseif c == "/" and source:sub(i+1, i+1) == "*" then
            i = i + 2
            while i <= len - 1 do
                if source:sub(i, i) == "*" and source:sub(i+1, i+1) == "/" then
                    i = i + 2
                    break
                end
                i = i + 1
            end
        else
            break
        end
    end
    return i
end

local function read_ident(source, i)
    if not is_alpha(source:sub(i, i)) then return nil, i end
    local a = i
    i = i + 1
    while i <= #source and is_alnum(source:sub(i, i)) do i = i + 1 end
    return source:sub(a, i - 1), i
end

local function read_balanced(source, open_char, close_char, i)
    if source:sub(i, i) ~= open_char then return nil, i end
    i = i + 1
    local len = #source
    local depth = 1
    local a = i
    while i <= len and depth > 0 do
        local c = source:sub(i, i)
        if c == open_char then
            depth = depth + 1
            i = i + 1
        elseif c == close_char then
            depth = depth - 1
            if depth == 0 then break end
            i = i + 1
        elseif c == '"' then
            i = i + 1
            while i <= len do
                if source:sub(i, i) == "\\" then i = i + 2
                elseif source:sub(i, i) == '"' then i = i + 1; break
                else i = i + 1 end
            end
        elseif c == "'" then
            i = i + 1
            while i <= len do
                if source:sub(i, i) == "\\" then i = i + 2
                elseif source:sub(i, i) == "'" then i = i + 1; break
                else i = i + 1 end
            end
        elseif c == "/" and source:sub(i+1, i+1) == "/" then
            while i <= len and source:sub(i, i) ~= "\n" do i = i + 1 end
        elseif c == "/" and source:sub(i+1, i+1) == "*" then
            i = i + 2
            while i <= len - 1 do
                if source:sub(i, i) == "*" and source:sub(i+1, i+1) == "/" then
                    i = i + 2
                    break
                end
                i = i + 1
            end
        else
            i = i + 1
        end
    end
    return source:sub(a, i - 1), i + 1
end

local function read_parens(source, i)  return read_balanced(source, "(", ")", i) end
local function read_braces(source, i)  return read_balanced(source, "{", "}", i) end
local function read_brackets(source, i) return read_balanced(source, "[", "]", i) end

local function scan_to_char(source, target, start)
    local len = #source
    local i = start
    while i <= len do
        local c = source:sub(i, i)
        if c == target then return i
        elseif c == "(" then local _, a = read_parens(source, i); i = a
        elseif c == "{" then local _, a = read_braces(source, i); i = a
        elseif c == "[" then local _, a = read_brackets(source, i); i = a
        elseif c == '"' then
            i = i + 1
            while i <= len do
                if source:sub(i, i) == "\\" then i = i + 2
                elseif source:sub(i, i) == '"' then i = i + 1; break
                else i = i + 1 end
            end
        elseif c == "'" then
            i = i + 1
            while i <= len do
                if source:sub(i, i) == "\\" then i = i + 2
                elseif source:sub(i, i) == "'" then i = i + 1; break
                else i = i + 1 end
            end
        elseif c == "/" and source:sub(i+1, i+1) == "/" then
            while i <= len and source:sub(i, i) ~= "\n" do i = i + 1 end
        elseif c == "/" and source:sub(i+1, i+1) == "*" then
            i = i + 2
            while i <= len - 1 do
                if source:sub(i, i) == "*" and source:sub(i+1, i+1) == "/" then
                    i = i + 2
                    break
                end
                i = i + 1
            end
        else
            i = i + 1
        end
    end
    return nil
end

-- ============================================================
-- Extract comma-separated identifier args from a parenthesized group
-- after a function-like macro call.
-- ============================================================

local function extract_ident_args(token, after_ident)
    local arg_start = skip_ws_and_comments(token, after_ident)
    if token:sub(arg_start, arg_start) ~= "(" then return {}, nil end
    local inner, after_paren = read_parens(token, arg_start)

    local args = {}
    local n = 1
    local len = #inner
    while n <= len do
        n = skip_ws_and_comments(inner, n)
        if n > len then break end
        local ident, after = read_ident(inner, n)
        if ident and ident ~= "" then
            table.insert(args, ident)
            n = after
        else
            n = n + 1
        end
        n = skip_ws_and_comments(inner, n)
        if n <= len and inner:sub(n, n) == "," then n = n + 1 end
    end

    return args, after_paren
end

-- ============================================================
-- Load WORD_COUNT manifest
-- ============================================================

local function load_word_counts(metadata_path)
    local counts = {}
    local content = read_file(metadata_path)
    local len = #content
    local i = 1
    local prefix = "WORD_COUNT("
    while i <= len do
        local nl = find_byte(content, "\n", i)
        local line_end = nl or (len + 1)
        local line = content:sub(i, line_end - 1)
        local trimmed = trim(line)
        if starts_with(trimmed, prefix) and ends_with(trimmed, ")") then
            local inner = trimmed:sub(#prefix + 1, #trimmed - 1)
            local comma = find_byte(inner, ",", 1)
            if comma then
                counts[trim(inner:sub(1, comma - 1))] =
                    tonumber(trim(inner:sub(comma + 1)))
            end
        end
        i = line_end + 1
    end
    return counts
end

-- ============================================================
-- Count words for a single comma-separated token
-- ============================================================

local function word_count_of_token(token, word_counts)
    local i = 1
    local len = #token
    while i <= len and is_space(token:sub(i, i)) do i = i + 1 end
    if i > len then return 0 end
    local name, after = read_ident(token, i)
    if not name then return 1 end
    local j = skip_ws_and_comments(token, after)
    if token:sub(j, j) == "(" then
        local wc = word_counts[name]
        if wc then return wc end
        io.stderr:write("  warning: unknown macro '" .. name .. "', assuming 1 word\n")
        return 1
    end
    return 1
end

-- ============================================================
-- Split brace-body into top-level comma-separated tokens
-- ============================================================

local function split_top_level_commas(body)
    local tokens = {}
    local len = #body
    local i = 1
    local token_start = 1
    while i <= len do
        local c = body:sub(i, i)
        if c == "(" then local _, a = read_parens(body, i); i = a
        elseif c == "{" then local _, a = read_braces(body, i); i = a
        elseif c == "[" then local _, a = read_brackets(body, i); i = a
        elseif c == '"' then
            i = i + 1
            while i <= len do
                if body:sub(i, i) == "\\" then i = i + 2
                elseif body:sub(i, i) == '"' then i = i + 1; break
                else i = i + 1 end
            end
        elseif c == "'" then
            i = i + 1
            while i <= len do
                if body:sub(i, i) == "\\" then i = i + 2
                elseif body:sub(i, i) == "'" then i = i + 1; break
                else i = i + 1 end
            end
        elseif c == "/" and body:sub(i+1, i+1) == "/" then
            while i <= len and body:sub(i, i) ~= "\n" do i = i + 1 end
        elseif c == "/" and body:sub(i+1, i+1) == "*" then
            i = i + 2
            while i <= len - 1 do
                if body:sub(i, i) == "*" and body:sub(i+1, i+1) == "/" then
                    i = i + 2
                    break
                end
                i = i + 1
            end
        elseif c == "," then
            table.insert(tokens, body:sub(token_start, i - 1))
            i = i + 1
            token_start = i
        else
            i = i + 1
        end
    end
    local last = body:sub(token_start, len)
    if trim(last) ~= "" then table.insert(tokens, last) end
    return tokens
end

-- ============================================================
-- Scan token for atom_label/atom_offset markers, walking through
-- balanced groups transparently (so nested calls are found)
-- ============================================================

local function scan_for_atom_markers(token, at_pos, labels, branches)
    local i = 1
    local len = #token
    while i <= len do
        i = skip_ws_and_comments(token, i)
        if i > len then break end
        local c = token:sub(i, i)
        if c == '"' then
            i = i + 1
            while i <= len do
                if token:sub(i, i) == "\\" then i = i + 2
                elseif token:sub(i, i) == '"' then i = i + 1; break
                else i = i + 1 end
            end
        elseif c == "'" then
            i = i + 1
            while i <= len do
                if token:sub(i, i) == "\\" then i = i + 2
                elseif token:sub(i, i) == "'" then i = i + 1; break
                else i = i + 1 end
            end
        elseif c == "/" and token:sub(i+1, i+1) == "/" then
            while i <= len and token:sub(i, i) ~= "\n" do i = i + 1 end
        elseif c == "/" and token:sub(i+1, i+1) == "*" then
            i = i + 2
            while i <= len - 1 do
                if token:sub(i, i) == "*" and token:sub(i+1, i+1) == "/" then
                    i = i + 2
                    break
                end
                i = i + 1
            end
        elseif is_alpha(c) then
            local ident, after = read_ident(token, i)
            if ident == "atom_label" then
                local args, after_paren = extract_ident_args(token, after)
                if #args >= 1 then labels[args[1]] = at_pos end
                if after_paren then i = after_paren else i = after end
            elseif ident == "atom_offset" then
                local args, after_paren = extract_ident_args(token, after)
                if #args >= 2 then
                    table.insert(branches, {
                        pos = at_pos,
                        target = args[2],
                        tag = args[1]
                    })
                end
                if after_paren then i = after_paren else i = after end
            else
                i = after
            end
        else
            i = i + 1
        end
    end
end

-- ============================================================
-- Scan atom body, count words, find markers
-- ============================================================

local function scan_atom_body(body, word_counts)
    local pos = 0
    local labels = {}
    local branches = {}

    for _, tok in ipairs(split_top_level_commas(body)) do
        local k = 1
        local tlen = #tok
        while k <= tlen and is_space(tok:sub(k, k)) do k = k + 1 end
        local leading_ident, leading_after = read_ident(tok, k)

        if leading_ident == "atom_label" or leading_ident == "atom_offset" then
            scan_for_atom_markers(tok, pos, labels, branches)
        else
            local words = word_count_of_token(tok, word_counts)
            scan_for_atom_markers(tok, pos, labels, branches)
            pos = pos + words
        end
    end

    return labels, branches, pos
end

-- ============================================================
-- Find every MipsAtom_(name) { ... } in a source
-- ============================================================

local function has_prefix(s, prefix)
    if #s < #prefix then return false end
    for i = 1, #prefix do
        if s:sub(i, i) ~= prefix:sub(i, i) then return false end
    end
    return true
end

local function skip_qualifiers(source, i)
    local keywords = {
        ["static"]=true, ["const"]=true, ["volatile"]=true,
        ["extern"]=true, ["register"]=true, ["auto"]=true,
        ["inline"]=true, ["typedef"]=true,
        ["internal"]=true, ["LP_"]=true, ["global"]=true, ["gkknown"]=true
    }
    while true do
        i = skip_ws_and_comments(source, i)
        local ident, after = read_ident(source, i)
        if not ident then return i end
        if keywords[ident] then i = after else return i end
    end
end

local function find_atoms(source_text)
    local atoms = {}
    local len = #source_text
    local i = 1

    local function try_wrapped(after_pos)
        local paren_pos = skip_ws_and_comments(source_text, after_pos)
        if source_text:sub(paren_pos, paren_pos) ~= "(" then return nil end
        local inner, after_paren = read_parens(source_text, paren_pos)
        local n = 1
        while n <= #inner and is_space(inner:sub(n, n)) do n = n + 1 end
        local ns = n
        while n <= #inner and is_alnum(inner:sub(n, n)) do n = n + 1 end
        local name = inner:sub(ns, n - 1)
        if name == "" then return nil end
        local brace_pos = scan_to_char(source_text, "{", after_paren)
        if not brace_pos then return nil end
        local body, after_brace = read_braces(source_text, brace_pos)
        return {name = name, body = body, after_brace = after_brace}
    end

    local function try_raw(after_pos)
        local next_pos = skip_ws_and_comments(source_text, after_pos)
        local next_ident, next_after = read_ident(source_text, next_pos)
        if not next_ident then return nil end
        if not has_prefix(next_ident, "code_") then return nil end
        if #next_ident <= 5 then return nil end
        local atom_name = next_ident:sub(6)
        local brace_pos = scan_to_char(source_text, "{", next_after)
        if not brace_pos then return nil end
        local body, after_brace = read_braces(source_text, brace_pos)
        return {name = atom_name, body = body, after_brace = after_brace}
    end

    while i <= len do
        i = skip_ws_and_comments(source_text, i)
        if i > len then break end
        i = skip_qualifiers(source_text, i)
        if i > len then break end

        local ident, after = read_ident(source_text, i)
        if not ident then
            i = i + 1
        elseif ident == "MipsAtom_" then
            local atom = try_wrapped(after)
            if atom then
                table.insert(atoms, {name = atom.name, body = atom.body})
                i = atom.after_brace
            else
                i = i + 1
            end
        elseif ident == "MipsCode" then
            local atom = try_raw(after)
            if atom then
                table.insert(atoms, {name = atom.name, body = atom.body})
                i = atom.after_brace
            else
                i = after
            end
        else
            i = after
        end
    end
    return atoms
end

-- ============================================================
-- Compute branch offsets (target - branch - 1)
-- ============================================================

local function compute_offsets(labels, branches)
    local results = {}
    for _, br in ipairs(branches) do
        local target = labels[br.target]
        if not target then
            error("Branch target '" .. br.target ..
                  "' has no atom_label (at word " .. br.pos .. ")")
        end
        table.insert(results, {
            target = br.target,
            tag = br.tag,
            offset = target - br.pos - 1
        })
    end
    return results
end

-- ============================================================
-- Generate header for one source
-- ============================================================

local function generate_header(source_path, atoms_data)
    local basename = basename_no_ext(source_path)
    local guard    = to_alnum_underscore(to_upper(basename)) .. "_OFFSETS_H"

    local lines = {}
    local function add(s) table.insert(lines, s) end

    add("// Auto-generated by gen_atom_offsets.lua — DO NOT EDIT")
    add("// Source: " .. source_path)
    add("#pragma once")
    add("")
    add("#pragma region " .. basename)
    add("")
    add("// Dispatch macro: token-pastes <tag>_<target> to the enum name")
    add("#undef atom_offset")
    add("#define atom_offset(tag, name)  atom_offset_##tag##_##name")
    add("")

    for _, atom in ipairs(atoms_data) do
        if #atom.offsets > 0 then
            add("// --- atom: " .. atom.name .. " (" .. atom.total_words .. " words) ---")
            add("")

            -- Build constant list once
            local consts = {}
            for _, r in ipairs(atom.offsets) do
                table.insert(consts, {
                    macro_name = "_atom_offset_" .. r.tag .. "_" .. r.target,
                    enum_name  = "atom_offset_" .. r.tag .. "_" .. r.target,
                    value      = r.offset
                })
            end

            -- Macro form: single source of truth for the literal value.
            -- Underscore prefix keeps it out of the C namespace so the
            -- enum can use the same identifier name without conflict.
            for _, c in ipairs(consts) do
                add("#define " .. pad_right(c.macro_name, 44) .. "  " .. c.value .. "")
            end
            add("")

            -- Enum form: C code uses the natural name, value comes
            -- from the underscore-prefixed macro so the literal is
            -- defined in exactly one place.
            add("enum {")
            for _, c in ipairs(consts) do
                add("    " .. c.enum_name .. " = " .. c.macro_name .. ",")
            end
            add("};")
            add("")
        end
    end

    add("#pragma endregion " .. basename)
    add("")
    return table.concat(lines, "\n") .. "\n"
end

-- ============================================================
-- Process one source
-- ============================================================

local function process_source(source_path, word_counts)
    local source = read_file(source_path)
    local atoms_raw = find_atoms(source)

    if #atoms_raw == 0 then
        io.stderr:write("  note: no MipsAtom_ declarations in " .. source_path .. "\n")
        return
    end

    local atoms_data = {}
    for _, atom in ipairs(atoms_raw) do
        local labels, branches, total = scan_atom_body(atom.body, word_counts)
        local offsets = compute_offsets(labels, branches)
        table.insert(atoms_data, {
            name = atom.name,
            total_words = total,
            offsets = offsets
        })
    end

    local basename = basename_no_ext(source_path)
    local out_dir = dirname(source_path) .. "/gen"
    ensure_dir(out_dir)
    local out_path = out_dir .. "/" .. basename .. ".offsets.h"
    write_file(out_path, generate_header(source_path, atoms_data))

    local total_branches = 0
    for _, a in ipairs(atoms_data) do total_branches = total_branches + #a.offsets end
    print("  " .. basename .. ": " .. #atoms_data .. " atom(s), " .. total_branches .. " branch(es)")
    for _, a in ipairs(atoms_data) do
        for _, r in ipairs(a.offsets) do
            print("    " .. a.name .. " -> " .. r.tag .. ":" .. r.target .. " : " .. r.offset)
        end
    end
end

-- ============================================================
-- Main
-- ============================================================

local function main(args)
    if #args < 2 then
        print("Usage: gen_atom_offsets.lua <metadata.h> <source1> [source2 ...]")
        os.exit(1)
    end
    local metadata_path = args[1]
    local sources = {}
    for i = 2, #args do table.insert(sources, args[i]) end
    local word_counts = load_word_counts(metadata_path)
    for _, src in ipairs(sources) do process_source(src, word_counts) end
end

main({...})