emu.speedmode("normal") -- Set the speed of the emulator

--[[

This script implements a simple "thread scheduler" for Super Mario Bros. 1-1.
It takes the idea that threads are "simply save states, combined with a condition for resuming them", and... well, proves it.

This demonstrates:
- Time slicing
    - Every 200 frames, we switch to a different thread
- Sleeping
    - When Mario kills an enemy, that thread will go to sleep for 300 frames
    - If all the threads are asleep, dead, or blocking, the game will crash!
- Critical Sections
    - When a Mario enters the "critical section area"", that thread enters a critical section, disabling the thread scheduler until he leaves the area
    - This models "when a thread is in a critical section, don't let anything else run until it's done"
- Mutexes
    - When a Mario enters the pipe, that thread attempts to gain a mutex, blocking other threads from entering the pipe until it leaves
    - This models "only one thread can be in this area at a time, but threads can run freely outside of it"
- Condition Variables
    - When a Mario touches the flagpole, it increments a condition variable, and waits for all threads to finish before continuing
    - This models "when a thread is waiting on a condition variable, it blocks until the condition is met"
    - Practically, what this looks like is that when Mario touches the flagpole, it instantly switches to another Mario, until all of them have touched it.
- Killing a thread
    - When a Mario dies, that thread is killed, and the thread scheduler switches to another thread
]]

-- Whether to draw the GUI on the screen
DRAW_GUI = true 

-- The number of threads we want to create
NUM_THREADS = 3

-- The number of threads which must reach the flagpole before the game continues
NUM_THREADS_MUST_REACH_FLAGPOLE = NUM_THREADS

-- The number of frames to sleep after killing an enemy
ENEMY_KILLED_SLEEP_TIME = 300

-- How often the scheduler should run (eg run every 20 frames)
SCHEDULER_FREQUENCY = 170

-- A "level layout" value for World 1-1 Main -- we use this to determine what level/sublevel Mario is in
LEVEL_LAYOUT_1_1 = 144

-- A "level layout" value for World 1-1 Pipe -- we use this to determine what level/sublevel Mario is in
LEVEL_LAYOUT_1_1_PIPE = 123

-- The X/Y bounds of the "critical section area" of World 1-1
-- This is the first area of the level, on top of the three blocks near where the first Goomba is
-- When Mario enters this area, the thread scheduler is disabled until he leaves
CRITICAL_AREA_X_START = 320
CRITICAL_AREA_X_END = 385
CRITICAL_AREA_Y_MAXIMUM = 125

-- A ConditionVariable marking how many Marios have finished (touched the flagpole)
-- We use this in a bit of a strange way:
-- 1. When a Mario touches the flagpole, we increment this condition variable
-- 2. Then, on the same thread, we wait for the condition variable to be equal to the number of threads (this blocks the thread)
-- 3. When the condition variable is equal to the number of threads, all threads will be unblocked
numMariosFinishedConditionVariable = {
    type = "ConditionVariable",
    value = 0,
    numWaiting = 0 -- The number of threads waiting on this condition variable
}

-- A Mutex that is used to ensure only one Mario is within the "mutex area" at a time
-- When a Mario enters the mutex area, it will attempt to gain the mutex, which will block the thread
pipeMutex = {
    type = "Mutex",
    owner = nil, -- The thread that currently owns the mutex
    numWaiting = 0 -- The number of threads waiting for the mutex
}

-- Thread states
THREAD_STATE_RUNNING = 'RUNNING' -- Thread is currently running
THREAD_STATE_PREEMPTED = 'PREEMPTED' -- Thread has been preempted, but is not blocked on anything 
THREAD_STATE_CRITICAL_SECTION = 'CRITICAL_SECTION' -- Thread is in a critical section and has disabled the thread scheduler entirely
THREAD_STATE_BLOCKED = 'BLOCKED' -- Thread is blocked on a condition variable or mutex
THREAD_STATE_SLEEPING = 'SLEEPING' -- Thread is sleeping. This could be combined into BLOCKED, but we keep it separate for clarity
THREAD_STATE_KILLED = 'KILLED' -- Thread has been killed

local threads = {}
local curFrame = 0
local lastSwitched = 0
local initiated = false -- Whether the player has started world 1-1 and we've started the threads

-- Returns the number of game frames until the next scheduler run
function untilNextScheduler()
    local delta = curFrame - lastSwitched
    if delta >= SCHEDULER_FREQUENCY then
        lastSwitched = curFrame
        return 0 -- Time to switch threads
    else
        return SCHEDULER_FREQUENCY - delta -- Time until next switch
    end
end

-- Returns the player's X position in the game
function playerX()
    local pagePosition = memory.readbyte(0x006D) * 256
    local absolutePosition = pagePosition + memory.readbyte(0x0086)
    return absolutePosition
end

-- Returns the player's Y position in the game
-- This increases the further *down* the screen the player is
function playerY()
    return memory.readbyte(0x00B5) * memory.readbyte(0x00CE)
end

-- Checks to see if the player has started world 1-1, and if so, initiates the script and launches threads
function initiate()
    emu.frameadvance()

    if emu.emulating() then
        local pos = playerX()
        local gameMode = memory.readbyte(0x0770)
        local preLevelScreen = memory.readbyte(0x0757) -- Read the pre-level screen value

        if gameMode == 1 and (preLevelScreen == 1 or pos ~= 40) then
            -- If the game is in the correct mode and the player is in the right position, we can start
            -- Create a number of save states, storing them in threads
            for i = 1, NUM_THREADS do
                local thread = {}
                thread.id = i
                thread.name = "Thread " .. i
                thread.state = THREAD_STATE_PREEMPTED -- Set the initial state of the thread
                thread.saveState = savestate.create()
                thread.color = (i == 1) and "red" or (i == 2) and "green" or "blue" -- Assign colors to threads
                thread.palette = (i == 1) and 0x3E or (i == 2) and 0x5E or 0x9E -- Assign palettes to threads
                thread.sleepingUntil = nil -- The frame until which this thread is sleeping, if it is in the SLEEPING state
                thread.blockedOn = nil -- The mutex or condition variable this thread is blocked on
                thread.desiredValueOfConditionVariable = nil -- The desired value of the condition variable this thread is waiting on

                thread.threadLocal = {} -- Thread-local storage for this thread
                thread.threadLocal.hasTouchedFlagpole = false -- Used to track if the thread has touched the flagpole
                thread.threadLocal.enemiesKilled = {false, false, false, false, false} -- Used to track which of the enemies we already know is dead

                savestate.save(thread.saveState) -- Save the current state

                table.insert(threads, thread)
            end

            initiated = true
            curThread = threads[1] -- Set the current thread to the first thread
            threadScheduler() -- Run the thread scheduler to select/start the first thread
        end
    end
end

-- Attempts to unblock a thread that is blocked on a condition variable or mutex
-- Returns true if the thread was successfully unblocked, false otherwise
function tryUnblock(thread)
    if thread.blockedOn.type == "ConditionVariable" and thread.blockedOn.value == thread.desiredValueOfConditionVariable then
        thread.blockedOn.numWaiting = thread.blockedOn.numWaiting - 1 -- Decrement the number of threads waiting on the condition variable
        thread.blockedOn = nil
        thread.state = THREAD_STATE_PREEMPTED
        return true
    elseif thread.blockedOn.type == "Mutex" then
        if pipeMutex.owner == nil then
            pipeMutex.owner = thread.id -- Set the current thread as the owner of the mutex
            thread.blockedOn.numWaiting = thread.blockedOn.numWaiting - 1 -- Decrement the number of threads waiting for the mutex
            thread.blockedOn = nil
            thread.state = THREAD_STATE_PREEMPTED
            return true
        elseif pipeMutex.owner == thread.id then
            -- If the thread already owns the mutex, it can continue
            thread.blockedOn = nil
            thread.state = THREAD_STATE_PREEMPTED
            return true
        else
            -- The mutex is owned by another thread, so we cannot unblock this thread
            return false
        end
    end

    return false
end

-- The core thread scheduler function
-- This function is responsible for saving the current thread's state, determining the next thread to run, and loading it in
function threadScheduler()
    local oldThreadNum = curThread.id
    local oldThread = curThread

    -- If current thread is in a critical section, do not switch threads
    if oldThread.state == THREAD_STATE_CRITICAL_SECTION then
        emu.print("Thread " .. oldThread.name .. " is in a critical section, skipping thread switch.")
        return -- If the current thread is in a critical section, do not switch threads
    end

    -- Save old context
    savestate.save(oldThread.saveState)

    -- If the old thread is running, set it to preempted state
    -- However, it could be in other states, such as BLOCKED or KILLED, meaning
    -- we don't want to overwrite that state -- just switch away from it
    if oldThread.state == THREAD_STATE_RUNNING then
        oldThread.state = THREAD_STATE_PREEMPTED -- Set the old thread to preempted state
    end

    -- Determine new thread
    for i = 0, NUM_THREADS - 2 do
        local newThreadNum = ((oldThreadNum + i) % NUM_THREADS) + 1 -- Cycle through threads, starting at the next thread from the current one
        local thread = threads[newThreadNum] -- Lua uses 1-based indexing

        emu.print("Current thread: " .. oldThreadNum .. ", New thread: " .. newThreadNum)
        emu.print("Checking thread " .. newThreadNum .. " with state: " .. thread.state)

        -- If a thread is PREEMPTED, that means it's not waiting on anything -- switch to it immediately
        if thread.state == THREAD_STATE_PREEMPTED then
            curThread = thread -- Set the current thread to the new thread
            emu.print("Switching to thread " .. newThreadNum .. " (PREEMPTED)")
            break -- Exit the loop once we find a thread to switch to
        end

        -- If a thread is SLEEPING, check if it can be woken up
        if thread.state == THREAD_STATE_SLEEPING then
            if curFrame >= thread.sleepingUntil then
                thread.state = THREAD_STATE_PREEMPTED -- Wake up the thread by setting it to preempted state
                curThread = thread -- Set the current thread to the new thread
                emu.print("Switching to thread " .. newThreadNum .. " (SLEEPING but can be woken up)")
                break -- Exit the loop once we find a thread to switch to
            else
                emu.print("Thread " .. newThreadNum .. " is SLEEPING, continuing search.")
            end
        end

        -- If a thread is BLOCKED, try to resolve the block
        if thread.state == THREAD_STATE_BLOCKED then
            if tryUnblock(thread) then
                curThread = thread -- Set the current thread to the new thread (adjusting for 0-based indexing)
                emu.print("Switching to thread " .. newThreadNum .. " (BLOCKED but can be unblocked)")
                break -- Exit the loop once we find a thread to switch to
            else
                emu.print("Thread " .. newThreadNum .. " is BLOCKED and cannot be unblocked, continuing search.")
            end
        end
    end

    -- HACK: Detect if there are no threads that can run
    --       We do this by checking if the thread we "just chose" is not in a runnable state
    if curThread.state ~= THREAD_STATE_PREEMPTED and curThread.state ~= THREAD_STATE_RUNNING then
        emu.print("No runnable threads found, deadlock?")
        curThread = nil -- Set current thread to nil to indicate no runnable threads
        return -- Exit the scheduler if no runnable threads are found
    end

    savestate.load(curThread.saveState)
    memory.writebyte(0x0779, curThread.palette) -- Set the palette for the new thread
    curThread.state = THREAD_STATE_RUNNING -- Set the new thread to running state

    lastSwitched = curFrame -- Update the last switched frame
end

-- Puts the current thread to sleep for a specified number of frames
-- This will set the thread state to SLEEPING and run the thread scheduler immediately
function sleep(numFrames)
    curThread.state = THREAD_STATE_SLEEPING -- Set the thread state to sleeping
    curThread.sleepingUntil = curFrame + numFrames -- Set the frame until which this thread is sleeping

    threadScheduler() -- Run the thread scheduler to switch to another thread
end

-- Kills the current thread, setting its state to KILLED and running the thread scheduler immediately
function killThread()
    -- If the thread is already killed, do nothing
    if curThread.state == THREAD_STATE_KILLED then
        return
    end

    -- Set the thread state to killed
    curThread.state = THREAD_STATE_KILLED
    curThread.blockedOn = nil -- Clear any blocked state

    -- TODO: Wipe any mutexes this thread owns (or don't -- good demonstration of dangling mutexes)

    threadScheduler() -- Run the thread scheduler to switch to another thread
end

-- Increment the condition variable that tracks how many Marios have finished (touched the flagpole)
function incrementConditionVariable()
    numMariosFinishedConditionVariable.value = numMariosFinishedConditionVariable.value + 1
end

-- Blocks the current thread until the condition variable (# of Marios who have reached flag)
-- is equal to the desired value (usually the number of threads)
function waitOnConditionVariable(desiredValue)
    -- Check if the condition variable is already met
    if numMariosFinishedConditionVariable.value == desiredValue then
        return
    end

    -- If thread is running, block it on the condition variable
    if curThread.state == THREAD_STATE_RUNNING then
        curThread.state = THREAD_STATE_BLOCKED -- Block the thread
        curThread.blockedOn = numMariosFinishedConditionVariable -- Set the condition variable it is waiting on
        curThread.desiredValueOfConditionVariable = desiredValue -- Set the desired value of the condition variable
        numMariosFinishedConditionVariable.numWaiting = numMariosFinishedConditionVariable.numWaiting + 1 -- Increment the number of threads waiting on the condition variable

        threadScheduler() -- We're now blocked -- run the thread scheduler to switch to another thread
    end
end

-- Attempts to gain ownership of the mutex which controls access to the Pipe sub-area of World 1-1
-- If another thread already owns the mutex, this thread will block until it can gain ownership
-- Returns true if the thread is now blocked on the mutex, false if it already owns it or can gain it
function gainMutex()
    -- If thread already owns mutex, don't do nothin'
    if pipeMutex.owner == curThread.id then
        return
    end

    -- If mutex is not owned, gain ownership
    if pipeMutex.owner == nil then
        pipeMutex.owner = curThread.id -- Set the current thread as the owner of the mutex
        return
    end

    -- Otherwise, mutex is owned by another thread
    -- Mark this thread as waiting for the mutex
    curThread.state = THREAD_STATE_BLOCKED -- Block the thread
    curThread.blockedOn = pipeMutex -- Set the mutex it is waiting on
    pipeMutex.numWaiting = pipeMutex.numWaiting + 1 -- Increment the number of threads waiting for the mutex

    threadScheduler() -- We're now blocked -- run the thread scheduler to switch to another thread
end

-- Attempts to release the mutex if the current thread owns it
-- Returns true if the mutex was successfully released, false if the thread does not own it
function releaseMutex()
    -- If the current thread owns the mutex, release it
    if pipeMutex.owner == curThread.id then
        pipeMutex.owner = nil -- Release the mutex
        
        -- If there are threads waiting for the mutex, we should immediately run the scheduler
        if pipeMutex.numWaiting > 0 then
            threadScheduler() -- Run the thread scheduler to switch to another thread
        end
    end
end

-- Enters a critical section, disabling the thread scheduler until the thread exits the critical section
-- Returns true if the thread successfully entered the critical section, false otherwise
function enterCriticalSection()
    if curThread.state == THREAD_STATE_RUNNING then
        curThread.state = THREAD_STATE_CRITICAL_SECTION -- Set the thread state to critical section
    end
end

-- Exits the critical section, allowing the thread scheduler to run again
-- Returns true if the thread successfully exited the critical section, false if it was not in one
function exitCriticalSection()
    if curThread.state == THREAD_STATE_CRITICAL_SECTION then
        curThread.state = THREAD_STATE_RUNNING -- Set the thread state back to running
    end
end

function onEnemyKilled()
    sleep(ENEMY_KILLED_SLEEP_TIME) -- Sleep for 140 frames after killing an enemy
end

-- Assess the current game frame to determine if any of the following are true:
-- 1. Mario just died
-- 2. Mario just entered or exited the mutex area (pipe) or critical area (first three-block elevated platform of 1-1, where first goomba spawns)
-- 3. Mario is touching the flagpole and needs to A) increment condition variable and then B) wait on the condition variable
-- 4. Mario has just killed an enemy, and is in desperate need of a lawyer
function assessGameFrame()
    local x = playerX() -- Get the player's X position
    local y = playerY() -- Get the player's Y position

    local levelLayoutNumber = memory.readbyte(0x00e7) -- Read the level layout number

    -- Determine if Mario is fucking dead
    if memory.readbyte(0x000E) == 0x0B then
        killThread()
        return -- If Mario is dead, we don't want to do anything else
    end

    -- Determine if Mario is within the critical area
    if x >= CRITICAL_AREA_X_START and x <= CRITICAL_AREA_X_END and y <= CRITICAL_AREA_Y_MAXIMUM then
        -- Mario IS within the critical area. Disable thread scheduling.
        enterCriticalSection()
    else
        -- Mario IS NOT within the critical area. Enable thread scheduling.
        exitCriticalSection()
    end

    -- Determine if Mario is in the "mutex area" (pipe) based on the level layout number
    if levelLayoutNumber == LEVEL_LAYOUT_1_1 then
        -- Not in the pipe, so we can release the mutex if we own it (will do nothing if we don't own it)
        releaseMutex()
    end
    if levelLayoutNumber == LEVEL_LAYOUT_1_1_PIPE then
        -- In the pipe, so we need to try to gain the mutex (will do nothing if we already own it)
        gainMutex()
        return
    end

    -- Determine if Mario is now touching the flagpole, and thus, must wait for the condition variable to be met
    if memory.readbyte(0x001D) == 3 then
        -- Mario IS touching the flagpole. 
        -- If this is the first frame he's touching it, we need to increment the condition variable
        if not curThread.threadLocal.hasTouchedFlagpole then
            curThread.threadLocal.hasTouchedFlagpole = true -- Mark that we've incremented the condition variable
            incrementConditionVariable()
        end

        -- Block this thread until the condition variable is met (all Marios have touched the flagpole)
        waitOnConditionVariable(NUM_THREADS) -- Wait on the condition variable for all Marios to finish
        return
    end

    -- Handle dead enemies
    local hasKilledEnemy = false -- Flag to track if we have killed an enemy this frame
    for i = 0, 5 do
        local enemyState = memory.readbyte(0x001E + i) -- Read the enemy state
        if enemyState == 0x04 and not curThread.threadLocal.enemiesKilled[i + 1] then
            -- If the enemy is dead and we haven't already marked it as killed, mark it as killed
            curThread.threadLocal.enemiesKilled[i + 1] = true
            hasKilledEnemy = true -- Set the flag to true
        elseif enemyState == 0x00 and curThread.threadLocal.enemiesKilled[i + 1] then
            -- If the enemy is not dead and we have marked it as killed, unmark it
            curThread.threadLocal.enemiesKilled[i + 1] = false
        end
    end

    -- If we killed an enemy this frame, handle it
    if hasKilledEnemy then
        onEnemyKilled()
    end
end

-- Display the current game / thread state on the screen
function drawGui()
    if not DRAW_GUI then
        return -- If we don't want to draw the GUI, exit early
    end

    local gameMode = memory.readbyte(0x0770)

    local color = curThread.color
    if curThread.state == THREAD_STATE_CRITICAL_SECTION then
        gui.drawtext(10, 10, "IN CRITICAL SECTION! THREAD SCHEDULER DISABLED", color) -- Display a message on the screen
    else
        gui.drawtext(10, 10, "Scheduler runs in: " .. untilNextScheduler(), color) -- Display a message on the screen
    end

    -- Print all thread states on the screen
    for i, t in ipairs(threads) do
        if t.state == THREAD_STATE_BLOCKED then
            gui.drawtext(10, 10 + (i * 10), t.name .. ": " .. t.state .. " on " .. t.blockedOn.type, t.color)
        else
            gui.drawtext(10, 10 + (i * 10), t.name .. ": " .. t.state, t.color)
        end
    end

    -- Print the condition variable value and the owner of the pipe mutex
    if numMariosFinishedConditionVariable.numWaiting > 0 then
        gui.drawtext(10, 20 + (NUM_THREADS * 10),
            "Condition Variable: " .. numMariosFinishedConditionVariable.value .. "/" .. NUM_THREADS .. " | " ..
                numMariosFinishedConditionVariable.numWaiting .. " threads waiting", color)
    else
        gui.drawtext(10, 20 + (NUM_THREADS * 10),
            "Condition Variable: " .. numMariosFinishedConditionVariable.value .. "/" .. NUM_THREADS, color)
    end

    if pipeMutex.numWaiting > 0 then
        gui.drawtext(10, 30 + (NUM_THREADS * 10),
            "Mutex Owner: " .. (pipeMutex.owner or "None") .. " | " .. pipeMutex.numWaiting .. " threads waiting", color)
    else
        gui.drawtext(10, 30 + (NUM_THREADS * 10), "Mutex Owner: " .. (pipeMutex.owner or "None"), color)
    end
end

-- Main loop of the script, which runs every frame AFTER the player has started world 1-1
-- Runs the thread scheduler 
function loop()
    -- If we have no current thread, then the thread scheduler has been unable to find a runnable thread
    -- This means we're fucked. We're either in a deadlock state, or all threads have been killed.
    -- So, we should display a message on the screen
    if curThread == nil then
        -- Pause emulator if not paused
        if not emu.paused() then
            emu.frameadvance() -- Advance a frame to clear any messages on the screen
            emu.pause()
        end

        gui.drawtext(10, 10, "NO THREAD CAN RUN -- DEADLOCK?")
        emu.frameadvance() -- Advance a frame to display the message
        return
    end

    -- We have a current thread, so we should run it
    emu.frameadvance()
    curFrame = curFrame + 1

    -- Call the thread scheduler if it's time to switch threads
    if untilNextScheduler() == 0 then
        threadScheduler()
    end

     -- Draw the GUI with the current thread and game state
    drawGui()
    
    -- Based on the current state of the game, do stuff (eg gain/release mutex, enter/exit critical section, etc)
    -- This is where all the "game-related logic" happens
    assessGameFrame()
end

-- Root loop of our Lua script
while true do
    -- Don't do anything until we detect that the player has started playing world 1-1.
    -- Once we detect this, we'll spawn our threads and start the scheduler
    if not initiated then
        initiate()
    else
        -- If we have initiated, we can perform the main loop
        loop()
    end
end