Coroutine Management [WIP]

--[[
	Author: NeverCast
	Last Edited: 2/11/12
	You're welcome to use this code however you like, but please leave my credit at the top. Thanks
--]]

-- Prior code writing hypothosis:
-- How to become top level,

-- Climb stack, Find single shell instance
-- Kill shell instance, override pullEvent
-- Override os.shutdown
-- Exit code
-- Wait for Bios to quit and hook that to fire up your own functions

-- Add the functions here that you want to run once the environment has been overridden
-- In this bios this was shell and rednet.run, I've removed rednet.run as an example
-- Any functions here will be top level and recieve all the events!
--[[local loadPool = {
		function()
			os.run( {}, "rom/programs/shell" )
		end
}
--]]

-- Feel free to set this to false if you don't want the text printed
local showMessages = false

if background then return end
term.clear()
term.setCursorPos(1,1)
local function sort(table)
local max=getfenv(2).table.maxn
        if #table==max(table) then return table end
        local i=1
        for j=1,max(table) do
                if table[j] and j~=i then
                        table[i]=table[j]
                        table[j]=nil
                        i=i+1
                elseif table[j] then
                        i=i+1
                end
        end
        return table
end

local function coroutineManage()
        local bgCo={}
        local bgW={}
        local bgP={}
        _G.background={}
        function _G.background.addCo(f,...)
                local argsy={...}
                if not type(f)=="function" then error("function expected, got "..type(f),2) end
                local i=#bgCo+1
                bgCo[i]=coroutine.create(function() f(unpack(argsy)) end)
                os.queueEvent("coroutine_create",i)
                return i
        end
        function _G.background.killCo(id)
                if not type(id)=="number" then error("number expected, got "..type(id),2) end
                if id==1 then error("Cannot kill shell",2) end
                if not bgCo[id] then return false, "No such coroutine" end
                local l= table.remove(bgCo,id)
                os.queueEvent("coroutine_kill",id)
                return l
        end
        function _G.background.pauseCo(id)
                if not type(id)=="number" then error("number expected, got "..type(id),2) end
                if not bgCo[id] then return false, "No such coroutine" end
                if not bgP[i] then os.queueEvent("coroutine_pause",id) end
                bgP[id]=true
                return true
        end
        function _G.background.resumeCo(id)
                if not type(id)==number then error("number expected, got "..type(id),2) end
                if not bgCo[id] then return false, "No such coroutine" end
                if bgP[id] then os.queueEvent("coroutine_resume",id) end
                bgP[id]=false
                return true
        end
        function _G.background.getCoroutines()
                return {unpack(bgCo)}
        end
        function _G.background.isPaused(i)
                return not bgCo[i] and nil or (bgP[i] or false)
        end
        _G.background.addCo(dofile,"rom/programs/shell")
--print("test")
        local e={}
        while true do
                bgCo=sort(bgCo)
                bgW=sort(bgW)
                if #bgCo==0 then os.shutdown() end
                for i=1,#bgCo do
                        local n=bgCo[i]
--print(n)
--print(i)
                        if n then
                                if (not (bgW[i] and bgW[i]~=e[1]) or e[1]=="terminate") and not bgP[i] then
                                        --print("Passing event to coroutine.")
                                        local ok,param=coroutine.resume(bgCo[i],unpack(e))
                                        --print(ok,":",param)
                                        --print(table.concat(e,":"))
                                        if not ok or coroutine.status(bgCo[i])=="dead" then
                                        if not ok then os.queueEvent("coroutine_error",i,err) else os.queueEvent("coroutine_death",i) end
                                                --print("Oh noes, coroutine #"..i.." has crashed!")
                                                --print(param)
                                                bgCo[i]=nil
                                                bgW[i]=nil
                                                bgP[i]=nil
                                                if i==1 then os.shutdown() end
                                        else
                                                bgW[i]=param or false
                                                --print(param)
                                        end
                                end
                        end
                end
                if #bgCo<#bgP+1 then for k,v in pairs(bgP) do if v then os.queueEvent("coroutine_resume",k) end end bgP={} end
                e={os.pullEventRaw()}
                --print("Recieved event")
        end
end

if os.myEnv then
	if showMessages then
		print"Injection completed successfully"
	end
	return
end

local prevState = {}

table.size = function(tbl)
	local count = 0
	for _,__ in pairs(tbl) do
		count = count + 1
	end
	return count
end

local function getParentShell()
	for i = 0, 5 do
		local env = getfenv(i)
		if table.size(env) == 1 then
			local i,v = next(env) -- Grab first
			if i == "shell" then
				return v
			end
		end
	end
	return nil
end

local function recover()
	-- Set flag
	os.myEnv = true
	-- Put back the trampled environment
	os.startTimer = prevState.startTimer
	os.shutdown = prevState.shutdown
	os.pullEvent = prevState.pullEvent
	-- Launch shell like nothing happened

	prevState = nil

	term.setCursorPos(1,1)
	term.clear()

	background={}
	_G.background=background

	if showMessages then
		-- Feel free to remove this line, It's just a proof thing
		print"Look at me, I load before shell does"
	end

	local ok, err = pcall(coroutineManage)
	os.shutdown()
end

os.stage = {}

-- Stages:
-- injected: Overriding os functions
-- shutdown: Shutdown program has been called
-- jstincse: Shell is doing it's just-in-case
-- bioswait: Bios is waiting for a key press
-- biosexit: Bios has quit
-- complete: fully injected


local function setStage(stage)
	os.stage.currentStage = stage
end

local function getStage()
	return os.stage.currentStage
end

local function _os_pullEvent(_filter)
	_filter = _filter or ""
	if _filter == "key" then
		setStage("bioswait")
		return "key", 0
	elseif _filter == "timer" then
		setStage("shutdown")
		return "timer", 0
	end
end

local function _replSleep(dur)
    local timer = prevState.startTimer( dur )
	repeat
		local sEvent, param = prevState.pullEvent( "timer" )
	until param == timer
end

local function _os_shutdown()
	if getStage() == "shutdown" then
		setStage("jstincse")
	elseif getStage() == "bioswait" then
		setStage("biosexit")
	end
	if getStage() == "biosexit" then
		recover()
	end
end

local function _os_startTimer(seconds)
	return 0
end

local function inject()
	prevState.startTimer = os.startTimer
	prevState.shutdown = os.shutdown
	prevState.pullEvent = os.pullEvent
	os.shutdown = _os_shutdown
	os.pullEvent = _os_pullEvent
	os.startTimer = _os_startTimer
	setStage("injected")
	local shellImpl = getParentShell()
	shellImpl.exit()
end

-- Start everything
inject()