---A+ promises in Lua. Made to behave similar to JavaScript. Based on:
---https://github.com/zserge/lua-promises (MIT license)
---https://github.com/SilverEzhik/Promise.lua (MPL-2.0 license)
---https://github.com/evaera/roblox-lua-promise (MIT license)
---
---Slighlty reworked to improve performance, stability and pass tests.
---
---In the future, I might include it for promises to be available out of the box,
---and start moving API to using those, but at the moment Lua Language Server can’t
---handle generics well enough for those promises to keep types, so it would seem
---a bit counterproductive.
---@diagnostic disable

local Promise = {}

local mt = {
  __index = {},
  __call = function (self, value) return self:resolve(value) end,
  __tostring = function (self)
    return 'Promise<%s>' % {self._state or 'pending'}
  end
}

---@class Promise
local mi = mt.__index

---@type fun(self: any): boolean
Promise.is = function (self) return getmetatable(self) == mt end

---Override this function if you want to track unhandled promise rejections differently.
---@type nil|fun(err: string|any)
Promise.onUnhandledRejection = function (err) ac.warn('uncaught promise failure:', err) end

local RESOLVING = const('resolving')
local REJECTING = const('rejecting')
local RESOLVED = const('fulfilled')
local REJECTED = const('rejected')

local function create(success, failure, parent)
	return setmetatable({success = success, failure = failure, parent = parent}, mt)
end

local failed = setmetatable({}, {__mode = 'k'})
local emptyFn = function () end

local function finish(self, ok)
	self._state = ok and RESOLVING or REJECTING
  if not ok and not self.parent and Promise.onUnhandledRejection then
    local k = self._value == nil and failed or self._value
    setTimeout(function ()
      if not failed[k] and Promise.onUnhandledRejection then
        Promise.onUnhandledRejection(k ~= failed and k or nil)
      else
        failed[k] = nil
      end
    end, 1)
  end
  if self.queue then
    for i = 1, #self.queue do
      if ok then
        self.queue[i]:resolve(self._value)
      else
        self.queue[i]:reject(self._value)
      end
    end
	end
	self._state = ok and RESOLVED or REJECTED
  if not ok and self._raiseErrors then
    error(self._value, math.huge)
  end
end

local function callbable(f)
  local t = type(f)
	if t == 'table' then
		local m = getmetatable(f)
		return m ~= nil and type(m.__call) == 'function'
	end
	return t == 'function'
end

local function solveDeferredValue(self, promiseCb)
  local called = false
  local ok, err = pcall(mi.next, self._value, function(v)
    if called then return end
    called = true
    self._value = v
    promiseCb(self, true) 
  end, function(v)
    if called then return end
    called = true
    self._value = v
    promiseCb(self, false)
  end)
  if not ok and not called then
    called = true
    self._value = err
    promiseCb(self, false)
  end
end

local function checkSelf(self, v)
  return rawequal(v, self) or self.parent and checkSelf(self.parent, v)
end

local function fire(self, ok)
  self._state = ok and RESOLVING or REJECTING
	if ok and Promise.is(self._value) and not rawequal(self, self._value) then
	  solveDeferredValue(self, fire)
  else
    if ok and callbable(self.success) then
      ok, self._value = pcall(self.success, self._value)
    elseif not ok and callbable(self.failure) then
      failed[self._value or failed] = true
      ok, self._value = pcall(self.failure, self._value)
      if ok then self._state = RESOLVING end
    end
    if Promise.is(self._value) and checkSelf(self, self._value) then
      ok, self._value = false, 'resolving promise with itself'
    end
    if ok and Promise.is(self._value) then
      solveDeferredValue(self, finish)
    else
      finish(self, ok)
    end
  end
end

local function resolveImpl(self, ok, value)
	if self._state == nil then
		self._value = value
		fire(self, ok)
	end
	return self
end

---Returns a function with a value bound that will resolve a Promise with that value when called.
---@param value any
---@return fun()
function mi:bind(value)
	return function () resolveImpl(self, true, value) end
end

---Resolve a Promise.
---@param value any
---@return self
function mi:resolve(value)
	return resolveImpl(self, true, value)
end

---Reject a Promise.
---@param value string|any
---@return self
function mi:reject(value)
	return resolveImpl(self, false, value)
end

---Resolve or reject a Promise.
---@param err string|any
---@param data any
---@return self
function mi:complete(err, data)
  return err and self:reject(err) or self:resolve(data)
end

---Wait on a promise within a coroutine. Raises an error if promise returns an error.
---@async
---@return any
function mi:await()
  if not coroutine.isyieldable() then
    error("Promise:await() must be called from a yieldable coroutine")
  end
  if self._state == RESOLVED then
    return self._value
  elseif self._state == REJECTED then
    error(self._value)
  end
  local co = coroutine.running()
  local next = create(function (obj)
    coroutine.resume(co, nil, obj)
  end, function (err)
    coroutine.resume(co, err)
  end, self)
  next._raiseErrors = true
  if not self.queue then self.queue = {} end
  table.insert(self.queue, next)
  local err, obj = coroutine.yield()
  if err then error(err) else return obj end
end

---Use this function to join promises together and more, should work the same as `.then()` in JavaScript.
---@param success nil|fun(value: any): any
---@param failure nil|fun(err: string|any): any
---@return Promise
function mi:next(success, failure)
  local next = create(success, failure, self)
  if self._state == RESOLVED then
    next:resolve(self._value)
  elseif self._state == REJECTED then
    next:reject(self._value)
  else
    if not self.queue then self.queue = {} end
    table.insert(self.queue, next)
  end
  return next
end

---Registers a callback only for the rejection of a promise.
---@param failure fun(err: string|any)
---@return Promise
function mi:catch(failure)
  return self:next(nil, failure)
end

---Returns a promise that will cast rejections to `nil`.
---@return Promise
function mi:ok()
  return self:next(nil, emptyFn)
end

---Registers a callback that will be called when the promise is settled.
---The resolved value cannot be modified from the callback.
---@param onFinally? fun(): any
---@return Promise
function mi:finally(onFinally)
  local function fn(value)
    if callbable(onFinally) then
      pcall(onFinally)
    end
    return value
  end
  return self:next(fn, fn)
end

---Create a new Promise.
---@param fn nil|fun(resolve: (fun(result: any)), reject: (fun(err: string|any)))
---@return Promise
function Promise.new(fn)
  local d = create()
  if callbable(fn) then
    local ok, err = pcall(fn, d, function (err) d:reject(err) end)
    if not ok then d:reject(err) end
  end
  return d
end

---Creates a promise resolved with the given value.
---@param value any
---@return Promise
function Promise.resolve(value)
  return Promise.new(function(resolve) resolve(value) end)
end

---Creates a promise rejected with the given error.
---@param err string
---@return Promise
function Promise.reject(err)
  return Promise.new(function(_, reject) reject(err) end)
end

---Create a new coroutine and return a Promise waiting for it to complete. Note: this coroutine
---will raise errors publicly unless rejections are handled.
---@param fn fun() @This function will be ran from a coroutine.
---@return Promise
function Promise.coroutine(fn)
  return Promise.new(function (r, j)
    coroutine.create(function ()
      local ok, ret = pcall(fn);
      if ok then
        r(ret)
      else
        j(ret)
      end
    end, 'es')
  end)
end

---Returns a promise that is resolved with an array of results when all of the provided promises resolve
---or rejected when any promise is rejected.
---@param promises Promise[]
---@return Promise
function Promise.all(promises)
  return Promise.new(function(resolve, reject)
    local remaining = #promises
    local results = table.new(remaining, 0)
    if remaining == 0 then
      resolve(results)
      return
    end
    for i, promise in ipairs(promises) do
      promise:next(function(value)
        results[i] = value
        remaining = remaining - 1
        if remaining == 0 then
          resolve(results)
        end
      end, reject)
    end
  end)
end

---Returns a promise that is resolved or rejected when any of the provided promises resolve or reject.
---@param promises Promise[]
---@return Promise
function Promise.race(promises)
  return Promise.new(function(resolve, reject)
    for _, promise in ipairs(promises) do
      promise:next(resolve, reject)
    end
  end)
end

---Returns a promise that is resolved when any of the provided promises resolve or rejected when all
---of the provided promises reject.
---@param promises Promise[]
---@return Promise
function Promise.any(promises)
  return Promise.new(function(resolve, reject)
    local errors = {}
    local remaining = #promises
    if remaining == 0 then
        reject('no promises were resolved')
        return
    end
    for i, promise in ipairs(promises) do
      promise:next(resolve, function(reason)
        errors[i] = reason
        remaining = remaining - 1
        if remaining == 0 then
          reject(errors)
        end
      end)
    end
  end)
end

---Returns a promise that is resolved with an array of results when all of the provided promises resolve or reject
---@param promises Promise[]
---@return Promise
function Promise.allSettled(promises)
  return Promise.new(function(resolve)
    local remaining = #promises
    local results = table.new(remaining, 0)
    if remaining == 0 then
      resolve(results)
      return
    end
    for i, promise in ipairs(promises) do
      promise:next(
        function(value) results[i] = { status = 'fulfilled', value = value } end,
        function(reason) results[i] = { status = 'rejected', reason = reason } end
      ):finally(function()
        remaining = remaining - 1
        if remaining == 0 then
          resolve(results)
        end
      end)
    end
  end)
end

---Take a function taking a callback (either `fun(args…, callback: fun(ret: any))` or `fun(args…, callback: fun(err: string, ret: any))`)
---and return a function which instead would take the same arguments (except for callback) but return a promise.
---@param fn function
---@param flags nil|'e'|'E'|'f'|'p'|'ef'|'Ef'|'ep'|'efp' @Flags defining how callback is working: `e` if `err` is passed to callback first, `E` if `err` is passed second, `f` if callback is the first argument to the call (if not set, callback will be passed after the arguments), `p` if  callback is called with multiple data arguments (in that case, they will be packed using `table.pack()` for the return). `E` and `p` are not compatible.
function Promise.promisify(fn, flags)
  local err = flags and (string.find(flags, 'e', nil, true) and 1 
    or string.find(flags, 'E', nil, true) and 2 or 0)
  local pack = (flags and string.find(flags, 'p', nil, true)) ~= nil
  local first = (flags and string.find(flags, 'f', nil, true)) ~= nil
  if err == 2 and pack then error('incompatible flags', 2) end
  return function(...)
    local r, p = Promise.new(), table.pack(...)
    local c = err == 1 
      and (pack 
        and function(err, ...) r:complete(err, table.pack(...)) end 
        or function(err, ret) r:complete(err, ret) end)
      or err == 2 and function(ret, err) r:complete(err, ret) end
      or (pack 
        and function(...) r:resolve(table.pack(...)) end 
        or function(ret) r:resolve(ret) end)
    if first then
      fn(c, table.unpack(p, 1, p.n))
    else
      p[p.n + 1] = c
      fn(table.unpack(p, 1, p.n + 1))
    end
    return r
  end
end

return Promise