-- utility module to make the Lua 5.1 standard libraries behave more like Lua 5.2

if _VERSION == "Lua 5.1" then

   bit32 = require("bit32")

   -- table that maps each running coroutine to the coroutine that resumed it
   -- this is used to build complete tracebacks when "coroutine-friendly" pcall
   -- is used.
   local pcall_previous
   local pcall_callOf
   local xpcall_running
   local coroutine_running = coroutine.running

   -- the most powerful getmetatable we can get (preferably from debug)
   local sudo_getmetatable = getmetatable
   local _error = error
   local _type = type

   if type(debug) == "table" then

      local _G, package = _G, package
      local debug_setfenv = debug.setfenv
      debug.setuservalue = function(obj, value)
         if _type(obj) ~= "userdata" then
            _error("bad argument #1 to 'setuservalue' (userdata expected, got "..
                  _type(obj)..")", 2)
         end
         if value == nil then value = _G end
         if _type(value) ~= "table" then
            _error("bad argument #2 to 'setuservalue' (table expected, got "..
                  _type(value)..")", 2)
         end
         return debug_setfenv(obj, value)
      end

      local debug_getfenv = debug.getfenv
      debug.getuservalue = function(obj)
         if _type(obj) ~= "userdata" then
            return nil
         else
            local v = debug_getfenv(obj)
            if v == _G or v == package then
               return nil
            end
            return v
         end
      end

      if _type(debug.getmetatable) == "function" then
         sudo_getmetatable = debug.getmetatable
      end

      local debug_setmetatable = debug.setmetatable
      if _type(debug_setmetatable) == "function" then
         debug.setmetatable = function(value, tab)
            debug_setmetatable(value, tab)
            return value
         end
      end

      local debug_getinfo = debug.getinfo
      local function calculate_trace_level(co, level)
         if level ~= nil then
            for out = 1, 1/0 do
               local info = (co==nil) and debug_getinfo(out, "") or debug_getinfo(co, out, "")
               if info == nil then
                  local max = out-1
                  if level <= max then
                     return level
                  end
                  return nil, level-max
               end
            end
         end
         return 1
      end

      pcall_previous = {}
      pcall_callOf = {}
      xpcall_running = {}
      local stack_pattern = "\nstack traceback:"
      local stack_replace = ""
      local debug_traceback = debug.traceback
      debug.traceback = function (co, msg, level)
         local lvl
         local nilmsg
         if _type(co) ~= "thread" then
            co, msg, level = coroutine_running(), co, msg
         end
         if msg == nil then
            msg = ""
            nilmsg = true
         elseif _type(msg) ~= "string" then
            return msg
         end
         if co == nil then
            msg = debug_traceback(msg, level or 1)
         else
            local xpco = xpcall_running[co]
            if xpco ~= nil then
               lvl, level = calculate_trace_level(xpco, level)
               if lvl then
                  msg = debug_traceback(xpco, msg, lvl)
               else
                  msg = msg..stack_pattern
               end
               lvl, level = calculate_trace_level(co, level)
               if lvl then
                  local trace = debug_traceback(co, "", lvl)
                  msg = msg..trace:gsub(stack_pattern, stack_replace)
               end
            else
               co = pcall_callOf[co] or co
               lvl, level = calculate_trace_level(co, level)
               if lvl then
                  msg = debug_traceback(co, msg, lvl)
               else
                  msg = msg..stack_pattern
               end
            end
            co = pcall_previous[co]
            while co ~= nil do
               lvl, level = calculate_trace_level(co, level)
               if lvl then
                  local trace = debug_traceback(co, "", lvl)
                  msg = msg..trace:gsub(stack_pattern, stack_replace)
               end
               co = pcall_previous[co]
            end
         end
         if nilmsg then
            msg = msg:gsub("^\n", "")
         end
         msg = msg:gsub("\n\t%(tail call%): %?", "\000")
         msg = msg:gsub("\n\t%.%.%.\n", "\001\n")
         msg = msg:gsub("\n\t%.%.%.$", "\001")
         msg = msg:gsub("(%z+)\001(%z+)", function(some, other)
            return "\n\t(..."..#some+#other.."+ tail call(s)...)"
         end)
         msg = msg:gsub("\001(%z+)", function(zeros)
            return "\n\t(..."..#zeros.."+ tail call(s)...)"
         end)
         msg = msg:gsub("(%z+)\001", function(zeros)
            return "\n\t(..."..#zeros.."+ tail call(s)...)"
         end)
         msg = msg:gsub("%z+", function(zeros)
            return "\n\t(..."..#zeros.." tail call(s)...)"
         end)
         msg = msg:gsub("\001", function(zeros)
            return "\n\t..."
         end)
         return msg
      end
   end

   local _pairs = pairs
   pairs = function(t)
      local mt = sudo_getmetatable(t)
      if _type(mt) == "table" and _type(mt.__pairs) == "function" then
         return mt.__pairs(t)
      else
         return _pairs(t)
      end
   end

   local _ipairs = ipairs
   ipairs = function(t)
      local mt = sudo_getmetatable(t)
      if _type(mt) == "table" and _type(mt.__ipairs) == "function" then
         return mt.__ipairs(t)
      else
         return _ipairs(t)
      end
   end

   local _setfenv = setfenv

   local function check_mode(mode, prefix)
      local has = { text = false, binary = false }
      for i = 1,#mode do
         local c = mode:sub(i, i)
         if c == "t" then has.text = true end
         if c == "b" then has.binary = true end
      end
      local t = prefix:sub(1, 1) == "\27" and "binary" or "text"
      if not has[t] then
         return "attempt to load a "..t.." chunk (mode is '"..mode.."')"
      end
   end

   local _load, _loadstring = load, loadstring
   load = function(ld, source, mode, env)
      mode = mode or "bt"
      local chunk, msg
      if _type( ld ) == "string" then
         if mode ~= "bt" then
            local merr = check_mode(mode, ld)
            if merr then return nil, merr end
         end
         chunk, msg = _loadstring(ld, source)
      else
         local ld_type = _type(ld)
         if ld_type ~= "function" then
            _error("bad argument #1 to 'load' (function expected, got "..ld_type..")", 2)
         end
         if mode ~= "bt" then
            local checked, merr = false, nil
            local function checked_ld()
               if checked then
                  return ld()
               else
                  checked = true
                  local v = ld()
                  merr = check_mode(mode, v or "")
                  if merr then return nil end
                  return v
               end
            end
            chunk, msg = _load(checked_ld, source)
            if merr then return nil, merr end
         else
            chunk, msg = _load(ld, source)
         end
      end
      if not chunk then
         return chunk, msg
      end
      if env ~= nil then
         _setfenv(chunk, env)
      end
      return chunk
   end

   loadstring = load

   local _loadfile = loadfile
   local io_open = io.open
   loadfile = function(file, mode, env)
      mode = mode or "bt"
      if mode ~= "bt" then
         local f = io_open(file, "rb")
         if f then
            local prefix = f:read(1)
            f:close()
            if prefix then
               local merr = check_mode(mode, prefix)
               if merr then return nil, merr end
            end
         end
      end
      local chunk, msg = _loadfile(file)
      if not chunk then
         return chunk, msg
      end
      if env ~= nil then
         _setfenv(chunk, env)
      end
      return chunk
   end

   function rawlen(v)
      local t = _type(v)
      if t ~= "string" and t ~= "table" then
         _error("bad argument #1 to 'rawlen' (table or string expected)", 2)
      end
      return #v
   end

   local gc_isrunning = true
   local _collectgarbage = collectgarbage
   local math_floor = math.floor
   collectgarbage = function(opt, ...)
      opt = opt or "collect"
      local v = 0
      if opt == "collect" then
         v = _collectgarbage(opt, ...)
         if not gc_isrunning then _collectgarbage("stop") end
      elseif opt == "stop" then
         gc_isrunning = false
         return _collectgarbage(opt, ...)
      elseif opt == "restart" then
         gc_isrunning = true
         return _collectgarbage(opt, ...)
      elseif opt == "count" then
         v = _collectgarbage(opt, ...)
         return v, (v-math_floor(v))*1024
      elseif opt == "step" then
         v = _collectgarbage(opt, ...)
         if not gc_isrunning then _collectgarbage("stop") end
      elseif opt == "isrunning" then
         return gc_isrunning
      elseif opt ~= "generational" and opt ~= "incremental" then
         return _collectgarbage(opt, ...)
      end
      return v
   end

   local os_execute = os.execute
   local bit32_rshift = bit32.rshift
   os.execute = function(cmd)
      local code = os_execute(cmd)
      -- Lua 5.1 does not report exit by signal.
      if code == 0 then
         return true, "exit", code
      else
         return nil, "exit", bit32_rshift(code, 8)
      end
   end

   local _select = select
   table.pack = function(...)
      return { n = _select('#', ...), ... }
   end

   table.unpack = unpack

   local main_coroutine = coroutine.create(function() end)

   local _assert = assert
   local _pcall = pcall
   local coroutine_create = coroutine.create
   coroutine.create = function (func)
      local success, result = _pcall(coroutine_create, func)
      if not success then
         _assert(_type(func) == "function", "bad argument #1 (function expected)")
         result = coroutine_create(function(...) return func(...) end)
      end
      return result
   end

   local pcall_mainOf = {}

   coroutine.running = function()
      local co = coroutine_running()
      if co then
         return pcall_mainOf[co] or co, false
      else
         return main_coroutine, true
      end
   end

   local coroutine_yield = coroutine.yield
   coroutine.yield = function(...)
      local co = coroutine_running()
      if co then
         return coroutine_yield(...)
      else
         _error("attempt to yield from outside a coroutine", 0)
      end
   end

   local coroutine_resume = coroutine.resume
   coroutine.resume = function(co, ...)
      if co == main_coroutine then
         return false, "cannot resume non-suspended coroutine"
      else
         return coroutine_resume(co, ...)
      end
   end

   local coroutine_status = coroutine.status
   coroutine.status = function(co)
      local notmain = coroutine_running()
      if co == main_coroutine then
         return notmain and "normal" or "running"
      else
         return coroutine_status(co)
      end
   end

   local function pcall_results(current, call, success, ...)
      if coroutine_status(call) == "suspended" then
         return pcall_results(current, call, coroutine_resume(call, coroutine_yield(...)))
      end
      if pcall_previous then
         pcall_previous[call] = nil
         local main = pcall_mainOf[call]
         if main == current then current = nil end
         pcall_callOf[main] = current
      end
      pcall_mainOf[call] = nil
      return success, ...
   end
   local function pcall_exec(current, call, ...)
      local main = pcall_mainOf[current] or current
      pcall_mainOf[call] = main
      if pcall_previous then
         pcall_previous[call] = current
         pcall_callOf[main] = call
      end
      return pcall_results(current, call, coroutine_resume(call, ...))
   end
   local coroutine_create52 = coroutine.create
   local function pcall_coroutine(func)
      if _type(func) ~= "function" then
         local callable = func
         func = function (...) return callable(...) end
      end
      return coroutine_create52(func)
   end
   pcall = function (func, ...)
      local current = coroutine_running()
      if not current then return _pcall(func, ...) end
      return pcall_exec(current, pcall_coroutine(func), ...)
   end

   local _tostring = tostring
   local function xpcall_catch(current, call, msgh, success, ...)
      if not success then
         xpcall_running[current] = call
         local ok, result = _pcall(msgh, ...)
         xpcall_running[current] = nil
         if not ok then
            return false, "error in error handling (".._tostring(result)..")"
         end
         return false, result
      end
      return true, ...
   end
   local _xpcall = xpcall
   local _unpack = unpack
   xpcall = function(f, msgh, ...)
      local current = coroutine_running()
      if not current then
         local args, n = { ... }, _select('#', ...)
         return _xpcall(function() return f(_unpack(args, 1, n)) end, msgh)
      end
      local call = pcall_coroutine(f)
      return xpcall_catch(current, call, msgh, pcall_exec(current, call, ...))
   end

   local math_log = math.log
   math.log = function(x, base)
      if base ~= nil then
         return math_log(x)/math_log(base)
      else
         return math_log(x)
      end
   end

   local table_concat = table.concat
   package.searchpath = function(name, path, sep, rep)
      sep = (sep or "."):gsub("(%p)", "%%%1")
      rep = (rep or package.config:sub(1, 1)):gsub("(%%)", "%%%1")
      local pname = name:gsub(sep, rep):gsub("(%%)", "%%%1")
      local msg = {}
      for subpath in path:gmatch("[^;]+") do
         local fpath = subpath:gsub("%?", pname)
         local f = io_open(fpath, "r")
         if f then
            f:close()
            return fpath
         end
         msg[#msg+1] = "\n\tno file '" .. fpath .. "'"
      end
      return nil, table_concat(msg)
   end

   local p_index = { searchers = package.loaders }
   local _rawset = rawset
   setmetatable(package, {
      __index = p_index,
      __newindex = function(p, k, v)
         if k == "searchers" then
            _rawset(p, "loaders", v)
            p_index.searchers = v
         else
            _rawset(p, k, v)
         end
      end
   })

   local string_gsub = string.gsub
   local function fix_pattern(pattern)
      return (string_gsub(pattern, "%z", "%%z"))
   end

   local string_find = string.find
   function string.find(s, pattern, ...)
      return string_find(s, fix_pattern(pattern), ...)
   end

   local string_gmatch = string.gmatch
   function string.gmatch(s, pattern)
      return string_gmatch(s, fix_pattern(pattern))
   end

   function string.gsub(s, pattern, ...)
      return string_gsub(s, fix_pattern(pattern), ...)
   end

   local string_match = string.match
   function string.match(s, pattern, ...)
      return string_match(s, fix_pattern(pattern), ...)
   end

   local string_rep = string.rep
   function string.rep(s, n, sep)
      if sep ~= nil and sep ~= "" and n >= 2 then
         return s .. string_rep(sep..s, n-1)
      else
         return string_rep(s, n)
      end
   end

   local string_format = string.format
   do
      local addqt = {
         ["\n"] = "\\\n",
         ["\\"] = "\\\\",
         ["\""] = "\\\""
      }

      local function addquoted(c)
         return addqt[c] or string_format("\\%03d", c:byte())
      end

      function string.format(fmt, ...)
         local args, n = { ... }, _select('#', ...)
         local i = 0
         local function adjust_fmt(lead, mods, kind)
            if #lead % 2 == 0 then
               i = i + 1
               if kind == "s" then
                  args[i] = _tostring(args[i])
               elseif kind == "q" then
                  args[i] = '"'..string_gsub(args[i], "[%z%c\\\"\n]", addquoted)..'"'
                  return lead.."%"..mods.."s"
               end
            end
         end
         fmt = string_gsub(fmt, "(%%*)%%([%d%.%-%+%# ]*)(%a)", adjust_fmt)
         return string_format(fmt, _unpack(args, 1, n))
      end
   end

   local io_write = io.write
   local io_output = io.output
   function io.write(...)
      local res, msg, errno = io_write(...)
      if res then
         return io_output()
      else
         return nil, msg, errno
      end
   end

   local lines_iterator
   do
      local function helper( st, var_1, ... )
         if var_1 == nil then
            if st.doclose then st.f:close() end
            if (...) ~= nil then
               _error((...), 2)
            end
         end
         return var_1, ...
      end

      function lines_iterator(st)
         return helper(st, st.f:read(_unpack(st, 1, st.n)))
      end
   end

   local valid_format = { ["*l"] = true, ["*n"] = true, ["*a"] = true }

   local io_input = io.input
   function io.lines(fname, ...)
      local doclose, file, msg
      if fname ~= nil then
         doclose, file, msg = true, io_open(fname, "r")
         if not file then _error(msg, 2) end
      else
         doclose, file = false, io_input()
      end
      local st = { f=file, doclose=doclose, n=_select('#', ...), ... }
      for i = 1, st.n do
         if _type(st[i]) ~= "number" and not valid_format[st[i]] then
           _error("bad argument #"..(i+1).." to 'for iterator' (invalid format)", 2)
         end
      end
      return lines_iterator, st
   end

   do
      local io_stdout = io.stdout
      local io_type = io.type
      local file_meta = sudo_getmetatable(io_stdout)
      if _type(file_meta) == "table" and _type(file_meta.__index) == "table" then
         local file_write = file_meta.__index.write
         file_meta.__index.write = function(self, ...)
            local res, msg, errno = file_write(self, ...)
            if res then
               return self
            else
               return nil, msg, errno
            end
         end

         file_meta.__index.lines = function(self, ...)
            if io_type(self) == "closed file" then
               _error("attempt to use a closed file", 2)
            end
            local st = { f=self, doclose=false, n=_select('#', ...), ... }
            for i = 1, st.n do
               if _type(st[i]) ~= "number" and not valid_format[st[i]] then
                  _error("bad argument #"..(i+1).." to 'for iterator' (invalid format)", 2)
               end
            end
            return lines_iterator, st
         end
      end
   end

end
