-- MAME Controls Menu Plugin with ViewSwitch-style Input Assignment
local exports = {}
exports.name = "controls"
exports.version = "0.9"
exports.description = "MAME Controls Display with ViewSwitch-style Input Assignment"
exports.license = "MIT"
exports.author = { name = "Custom" }

function exports.startplugin()
    -- Variables to track key state
    local single_key_pressed = false
    local hotkey1_pressed = false
    local hotkey2_pressed = false
    
    -- Cooldown variables to prevent immediate re-triggering
    local last_show_time = 0
    local cooldown_period = 1.0  -- 1 second cooldown
    local showing_controls = false
    
    -- Get the plugin directory for config file
    local function get_config_path()
        local possible_paths = {
            "plugins/controls/controls_config.cfg",
            "preview/controls_config.cfg",
            "controls_config.cfg"
        }
        
        for _, path in ipairs(possible_paths) do
            local test_file = io.open(path, "a")
            if test_file then
                test_file:close()
                return path
            end
        end
        
        return "controls_config.cfg"
    end
    
    local config_file_path = get_config_path()
    
    -- Default configuration
    local default_config = {
        single_assign = nil,
        hotkey1 = nil,
        hotkey2 = nil,
        show_on_pause = false,
        lightning_mode = false
    }
    
    local config = {}
    
    -- Input assignment state (ViewSwitch-style polling)
    local input_poll = nil  -- Will hold polling state when active
    local commonui = nil    -- Will be loaded when needed
    
    -- Menu state tracking
    local menu_items = {}
    local menu_indices = {
        single_assign = nil,
        hotkey1 = nil,
        hotkey2 = nil,
        show_on_pause = nil,
        lightning_mode = nil,
        clear_single = nil,
        clear_hotkey1 = nil,
        clear_hotkey2 = nil,
        clear_all = nil
    }
    
    -- Menu state
    local menu_state = "main"
    
    -- Precaching variables
    local user_paused = false
    local current_rom = nil
    local rom_precached = false
    local has_timer = (emu.timer ~= nil)
    
    print("Controls plugin initializing...")
    print("Timer functionality available: " .. tostring(has_timer))
    
    -- Function to save configuration to file
    local function save_config()
        local file = io.open(config_file_path, "w")
        if file then
            file:write("-- MAME Controls Plugin Configuration\n")
            file:write("-- This file is automatically generated\n")
            file:write("-- Set to 'nil' or remove line to disable that assignment\n")
            file:write("\n")
            file:write("single_assign=" .. (config.single_assign or "nil") .. "\n")
            file:write("hotkey1=" .. (config.hotkey1 or "nil") .. "\n")
            file:write("hotkey2=" .. (config.hotkey2 or "nil") .. "\n")
            file:write("show_on_pause=" .. tostring(config.show_on_pause or false) .. "\n")
            file:write("lightning_mode=" .. tostring(config.lightning_mode or false) .. "\n")
            file:close()
            print("Configuration saved to: " .. config_file_path)
            return true
        else
            print("Error: Could not save configuration to: " .. config_file_path)
            return false
        end
    end
    
    -- Function to load configuration from file
    local function load_config()
        for key, value in pairs(default_config) do
            config[key] = value
        end
        
        local file = io.open(config_file_path, "r")
        if file then
            print("Loading configuration from: " .. config_file_path)
            for line in file:lines() do
                if not line:match("^%s*%-%-") and not line:match("^%s*$") then
                    local key, value = line:match("^([^=]+)=(.+)$")
                    if key and value then
                        key = key:gsub("%s+", "")
                        value = value:gsub("%s+", "")
                        
                        if value ~= "nil" then
                            if key == "show_on_pause" or key == "lightning_mode" then
                                config[key] = (value:lower() == "true")
                            else
                                config[key] = value
                            end
                            print("Loaded config: " .. key .. " = " .. tostring(config[key]))
                        end
                    end
                end
            end
            file:close()
        else
            print("No configuration file found, using defaults")
            save_config()
        end
    end
    
    -- Load configuration on startup
    load_config()
    
    -- Function to convert input code to readable name (improved)
    local function format_input_name(input_code)
        if not input_code then return "Not Set" end
        
        -- Handle keyboard keys
        if input_code:match("KEYCODE_") then
            local key = input_code:gsub("KEYCODE_", "")
            return "Key: " .. key
        end
        
        -- Handle gamepad inputs
        if input_code:match("JOYCODE_") then
            local parts = {}
            for part in input_code:gmatch("[^_]+") do
                table.insert(parts, part)
            end
            
            if #parts >= 3 then
                local player = parts[2] or "1"
                local button_part = parts[3]
                
                -- Handle numbered buttons
                if button_part:match("BUTTON(%d+)") then
                    local num = button_part:match("BUTTON(%d+)")
                    local button_names = {
                        ["1"] = "Button A",
                        ["2"] = "Button B", 
                        ["3"] = "Button X",
                        ["4"] = "Button Y",
                        ["5"] = "Left Bumper (LB)",
                        ["6"] = "Right Bumper (RB)",
                        ["7"] = "Left Trigger (LT)",
                        ["8"] = "Right Trigger (RT)",
                        ["9"] = "Left Stick Click (L3)",
                        ["10"] = "Right Stick Click (R3)",
                        ["11"] = "Start",
                        ["12"] = "Back/Select"
                    }
                    
                    local readable_name = button_names[num] or ("Button " .. num)
                    return "P" .. player .. " " .. readable_name
                    
                -- Handle named buttons
                elseif button_part == "START" then
                    return "P" .. player .. " Start"
                elseif button_part == "SELECT" then
                    return "P" .. player .. " Select/Back"
                elseif button_part:match("HAT") then
                    local direction = button_part:gsub("HAT1?", ""):gsub("UP", "Up"):gsub("DOWN", "Down"):gsub("LEFT", "Left"):gsub("RIGHT", "Right")
                    return "P" .. player .. " D-Pad " .. direction
                end
            end
        end
        
        -- Fallback: return the code as-is but cleaned up
        return input_code:gsub("JOYCODE_1_", ""):gsub("KEYCODE_", "")
    end
    
    -- Function to get input sequence name (like ViewSwitch does)
    local function get_input_sequence_name(input_code)
        if not input_code or not manager or not manager.machine or not manager.machine.input then
            return "None"
        end
        
        local input = manager.machine.input
        if not input.seq_from_tokens or not input.seq_name then
            return format_input_name(input_code)
        end
        
        local seq = input:seq_from_tokens(input_code)
        if seq then
            return input:seq_name(seq)
        else
            return format_input_name(input_code)
        end
    end
    
    -- Precaching function
    local function precache_controls(game_name)
        if not game_name or game_name == "" or game_name == "___empty" then
            print("No valid ROM name available, skipping precache")
            return 
        end
        
        if current_rom == game_name and rom_precached then 
            print("ROM already precached: " .. game_name)
            return 
        end
        
        os.execute("echo Game detected: " .. game_name .. " > C:\\mame_plugin_log.txt")
        
        local command = string.format('"preview\\mame controls.exe" --precache --game %s', game_name)
        print("Pre-caching controls for: " .. game_name)
        
        local result = os.execute(command)
        if result then
            current_rom = game_name
            rom_precached = true
            print("Precaching complete for: " .. game_name)
        else
            print("Precaching failed for: " .. game_name)
        end
    end

    -- Function to show controls
    local function show_controls()
        local current_time = os.clock()
        
        if current_time - last_show_time < cooldown_period then
            print("Controls cooldown active, ignoring request")
            return
        end
        
        if showing_controls then
            print("Controls already showing, ignoring request")
            return
        end
        
        local game_name = emu.romname()

        if game_name and game_name ~= "" and game_name ~= "___empty" then
            print("Showing controls for: " .. game_name .. " (Lightning Mode: " .. tostring(config.lightning_mode) .. ")")
            showing_controls = true
            last_show_time = current_time
            
            if emu.pause then
                emu.pause()
            end
            
            local command
            if config.lightning_mode then
                command = string.format('"preview\\mame controls.exe" --image-only --game %s', game_name)
                print("Using Lightning Mode (fast image-only)")
            else
                command = string.format('"preview\\mame controls.exe" --preview-only --game %s --clean-preview --screen 1', game_name)
                print("Using Normal Mode (full preview)")
            end
            
            print("Running: " .. command)
            os.execute(command)
            
            showing_controls = false
            
            if user_paused and emu.unpause then
                print("Unpausing MAME after controls")
                emu.unpause()
                user_paused = false
            elseif emu.unpause then
                emu.unpause()
            end
        else
            print("No valid ROM name available, cannot show controls")
        end
    end

    -- Menu population function (ViewSwitch-style)
    local function menu_populate()
        menu_items = {}
        
        -- Reset menu indices
        for key in pairs(menu_indices) do
            menu_indices[key] = nil
        end
        
        local game_name = emu.romname()
        
        if menu_state == "main" then
            if game_name and game_name ~= "" and game_name ~= "___empty" then
                table.insert(menu_items, {"Show Controls for " .. game_name, "", ""})
                table.insert(menu_items, {"Configure Hotkeys", "", ""})
                table.insert(menu_items, {"---", "", ""})
                table.insert(menu_items, {"Exit Controls: Any Keyboard or XINPUT Key", "", "off"})
            else
                table.insert(menu_items, {"Show Controls (No ROM loaded)", "", ""})
                table.insert(menu_items, {"Configure Hotkeys", "", ""})
                table.insert(menu_items, {"---", "", ""})
                table.insert(menu_items, {"Exit Controls: Any Keyboard or XINPUT Key", "", "off"})
            end
        elseif menu_state == "config" then
            table.insert(menu_items, {"← Back to Main Menu", "", ""})
            table.insert(menu_items, {"---", "", ""})
            
            -- Input assignment entries with current values and polling state
            local single_seq = "None"
            local single_flags = ""
            if config.single_assign then
                single_seq = get_input_sequence_name(config.single_assign)
            end
            if input_poll and input_poll.type == "single" then
                single_flags = "lr"
            end
            table.insert(menu_items, {"Single Key/Button", single_seq, single_flags})
            menu_indices.single_assign = #menu_items
            
            local hotkey1_seq = "None"
            local hotkey1_flags = ""
            if config.hotkey1 then
                hotkey1_seq = get_input_sequence_name(config.hotkey1)
            end
            if input_poll and input_poll.type == "hotkey1" then
                hotkey1_flags = "lr"
            end
            table.insert(menu_items, {"Hotkey 1", hotkey1_seq, hotkey1_flags})
            menu_indices.hotkey1 = #menu_items
            
            local hotkey2_seq = "None"
            local hotkey2_flags = ""
            if config.hotkey2 then
                hotkey2_seq = get_input_sequence_name(config.hotkey2)
            end
            if input_poll and input_poll.type == "hotkey2" then
                hotkey2_flags = "lr"
            end
            table.insert(menu_items, {"Hotkey 2", hotkey2_seq, hotkey2_flags})
            menu_indices.hotkey2 = #menu_items
            
            table.insert(menu_items, {"---", "", ""})
            
            -- Toggle settings
            table.insert(menu_items, {"Show on Pause", config.show_on_pause and "Enabled" or "Disabled", "lr"})
            menu_indices.show_on_pause = #menu_items
            
            table.insert(menu_items, {"Lightning Mode", config.lightning_mode and "Enabled" or "Disabled", "lr"})
            menu_indices.lightning_mode = #menu_items
            
            table.insert(menu_items, {"---", "", ""})
            
            -- Clear options
            table.insert(menu_items, {"Clear Single Assignment", "", ""})
            menu_indices.clear_single = #menu_items
            
            table.insert(menu_items, {"Clear Hotkey 1", "", ""})
            menu_indices.clear_hotkey1 = #menu_items
            
            table.insert(menu_items, {"Clear Hotkey 2", "", ""})
            menu_indices.clear_hotkey2 = #menu_items
            
            table.insert(menu_items, {"Clear All Assignments", "", ""})
            menu_indices.clear_all = #menu_items
            
            table.insert(menu_items, {"---", "", ""})
            
            -- Help text
            table.insert(menu_items, {"How to use:", "", "off"})
            table.insert(menu_items, {"• Single: Press one key/button to show controls", "", "off"})
            table.insert(menu_items, {"• Hotkeys: Press both buttons together", "", "off"})
            table.insert(menu_items, {"• Lightning Mode: Fast image-only display", "", "off"})
            table.insert(menu_items, {"• Normal Mode: Full preview with options", "", "off"})
            table.insert(menu_items, {"• Press UI_CLEAR to clear hotkey during assignment", "", "off"})
        end
        
        -- Apply polling overlay if active (ViewSwitch-style)
        if input_poll then
            return input_poll.poller:overlay(menu_items)
        else
            return menu_items
        end
    end

    -- Menu callback function (ViewSwitch-style input assignment)
    local function menu_callback(index, event)
        -- Handle polling state first (ViewSwitch pattern)
        if input_poll then
            -- Check if polling completed
            if input_poll.poller:poll() then
                if input_poll.poller.sequence then
                    -- Input was assigned
                    local input_code = manager.machine.input:seq_to_tokens(input_poll.poller.sequence)
                    if input_poll.type == "single" then
                        config.single_assign = input_code
                        print("Single key/button assigned: " .. format_input_name(input_code))
                    elseif input_poll.type == "hotkey1" then
                        config.hotkey1 = input_code
                        print("Hotkey 1 assigned: " .. format_input_name(input_code))
                    elseif input_poll.type == "hotkey2" then
                        config.hotkey2 = input_code
                        print("Hotkey 2 assigned: " .. format_input_name(input_code))
                    end
                    save_config()
                end
                input_poll = nil
                return true  -- Refresh menu
            end
            return false  -- Still polling
        end
        
        if event == "select" then
            if menu_state == "main" then
                if index == 1 then
                    -- Show controls
                    local game_name = emu.romname()
                    if game_name and game_name ~= "" and game_name ~= "___empty" then
                        show_controls()
                        return false  -- Don't refresh menu for this action
                    else
                        print("No ROM is currently loaded, cannot show controls")
                        return false
                    end
                elseif index == 2 then
                    -- Configure Settings
                    menu_state = "config"
                    return true  -- Refresh to show new menu
                end
            elseif menu_state == "config" then
                if index == 1 then
                    -- Back to main menu
                    menu_state = "main"
                    input_poll = nil  -- Cancel any active polling
                    return true  -- Refresh to show main menu
                elseif index == menu_indices.single_assign then
                    -- Assign Single Key/Button
                    if not commonui then
                        commonui = require('commonui')
                    end
                    input_poll = { type = "single", poller = commonui.switch_polling_helper() }
                    print("Waiting for single key/button input...")
                    return true  -- Refresh to show polling state
                elseif index == menu_indices.hotkey1 then
                    -- Assign Hotkey 1
                    if not commonui then
                        commonui = require('commonui')
                    end
                    input_poll = { type = "hotkey1", poller = commonui.switch_polling_helper() }
                    print("Waiting for hotkey 1 input...")
                    return true  -- Refresh to show polling state
                elseif index == menu_indices.hotkey2 then
                    -- Assign Hotkey 2
                    if not commonui then
                        commonui = require('commonui')
                    end
                    input_poll = { type = "hotkey2", poller = commonui.switch_polling_helper() }
                    print("Waiting for hotkey 2 input...")
                    return true  -- Refresh to show polling state
                elseif index == menu_indices.show_on_pause then
                    -- Toggle Show on Pause
                    config.show_on_pause = not config.show_on_pause
                    print("Show on Pause: " .. (config.show_on_pause and "Enabled" or "Disabled"))
                    save_config()
                    return true  -- Refresh to show updated setting
                elseif index == menu_indices.lightning_mode then
                    -- Toggle Lightning Mode
                    config.lightning_mode = not config.lightning_mode
                    print("Lightning Mode: " .. (config.lightning_mode and "Enabled" or "Disabled"))
                    save_config()
                    return true  -- Refresh to show updated setting
                elseif index == menu_indices.clear_single then
                    -- Clear Single Assignment
                    config.single_assign = nil
                    print("Single assignment cleared")
                    save_config()
                    return true  -- Refresh to show cleared assignment
                elseif index == menu_indices.clear_hotkey1 then
                    -- Clear Hotkey 1
                    config.hotkey1 = nil
                    print("Hotkey 1 cleared")
                    save_config()
                    return true  -- Refresh to show cleared assignment
                elseif index == menu_indices.clear_hotkey2 then
                    -- Clear Hotkey 2
                    config.hotkey2 = nil
                    print("Hotkey 2 cleared")
                    save_config()
                    return true  -- Refresh to show cleared assignment
                elseif index == menu_indices.clear_all then
                    -- Clear All Assignments
                    config.single_assign = nil
                    config.hotkey1 = nil
                    config.hotkey2 = nil
                    print("All assignments cleared")
                    save_config()
                    return true  -- Refresh to show all cleared assignments
                end
            end
        elseif event == "clear" then
            -- Handle clear events during input assignment (ViewSwitch pattern)
            if menu_state == "config" then
                if index == menu_indices.single_assign then
                    config.single_assign = nil
                    print("Single assignment cleared")
                    save_config()
                    return true
                elseif index == menu_indices.hotkey1 then
                    config.hotkey1 = nil
                    print("Hotkey 1 cleared")
                    save_config()
                    return true
                elseif index == menu_indices.hotkey2 then
                    config.hotkey2 = nil
                    print("Hotkey 2 cleared")
                    save_config()
                    return true
                end
            end
        end
        return false
    end

    -- Frame callback for hotkey checking
    emu.register_frame_done(function()        
        if not manager or not manager.machine then
            return false
        end
        
        local machine = manager.machine
        if not machine.input then
            return false
        end
        
        -- Skip hotkey detection if we're in polling mode
        if input_poll then
            return false
        end
        
        -- Regular hotkey detection
        local input = machine.input
        
        -- Check Single Assignment (if configured)
        if config.single_assign and input.seq_pressed then
            local single_key_state = false
            local seq = input:seq_from_tokens(config.single_assign)
            if seq then
                single_key_state = input:seq_pressed(seq)
            end
            
            if single_key_state and not single_key_pressed then
                print("Single assignment activated: " .. config.single_assign)
                show_controls()
            end
            single_key_pressed = single_key_state
        end
        
        -- Check Hotkey Combination (if both configured)
        if config.hotkey1 and config.hotkey2 and input.seq_pressed then
            local hotkey1_state = false
            local hotkey2_state = false
            
            local seq1 = input:seq_from_tokens(config.hotkey1)
            if seq1 then hotkey1_state = input:seq_pressed(seq1) end
            
            local seq2 = input:seq_from_tokens(config.hotkey2)
            if seq2 then hotkey2_state = input:seq_pressed(seq2) end
            
            if hotkey1_state and hotkey2_state and (not hotkey1_pressed or not hotkey2_pressed) then
                print("Hotkey combo activated: " .. config.hotkey1 .. " + " .. config.hotkey2)
                show_controls()
            end
            
            hotkey1_pressed = hotkey1_state
            hotkey2_pressed = hotkey2_state
        end
        
        return false
    end)

    -- Register menu
    emu.register_menu(menu_callback, menu_populate, "Controls")
    
    -- Register pause handler
    if emu.register_pause then
        print("Registering pause handler")
        emu.register_pause(function()
            local game_name = emu.romname()
            if not game_name or game_name == "" or game_name == "___empty" then
                print("No ROM is loaded, skipping pause handler")
                return
            end
            
            if not user_paused then
                user_paused = true
                print("User paused MAME")
                
                if config.show_on_pause then
                    print("Show on pause enabled - displaying controls")
                    show_controls()
                else
                    print("Show on pause disabled - not displaying controls")
                end
            else
                user_paused = false
                print("MAME unpaused")
            end
        end)
    else
        print("emu.register_pause not available in this MAME version")
    end
    
    -- Register start handler
    if emu.register_start then
        print("Registering start handler")
        emu.register_start(function()
            rom_precached = false
            
            if not has_timer then
                local game_name = emu.romname()
                if game_name and game_name ~= "" and game_name ~= "___empty" then
                    print("ROM detected via start handler (without timer): " .. game_name)
                    precache_controls(game_name)
                else
                    print("No valid ROM name available in start handler")
                end
                return
            end
            
            emu.timer.pulse(1.0, function()
                local game_name = emu.romname()
                if game_name and game_name ~= "" and game_name ~= "___empty" then
                    print("ROM detected via start handler: " .. game_name)
                    precache_controls(game_name)
                else
                    print("No valid ROM name available in start handler")
                end
                return false
            end)
        end)
    end
    
    -- Register stop handler
    emu.register_stop(function()
        user_paused = false
        rom_precached = false
        current_rom = nil
        showing_controls = false
        last_show_time = 0
        input_poll = nil  -- Clear any active polling
        print("Controls plugin reset for next game")
    end)
    
    -- Periodic timer fallback
    if has_timer then
        local periodic_timer = emu.timer.periodic(5.0, function()
            local game_name = emu.romname()
            if game_name and game_name ~= "" and game_name ~= "___empty" and game_name ~= current_rom then
                print("ROM change detected via periodic timer: " .. game_name)
                precache_controls(game_name)
            end
            return true
        end)
        print("Periodic timer registered")
    else
        print("Timer functionality not available, skipping periodic checks")
    end
    
    print("Controls plugin loaded with ViewSwitch-style input assignment")
    print("Single Assignment: " .. (config.single_assign and format_input_name(config.single_assign) or "None"))
    print("Hotkey 1: " .. (config.hotkey1 and format_input_name(config.hotkey1) or "None"))
    print("Hotkey 2: " .. (config.hotkey2 and format_input_name(config.hotkey2) or "None"))
    print("Show on Pause: " .. tostring(config.show_on_pause))
    print("Lightning Mode: " .. tostring(config.lightning_mode))
end

return exports