local ds = require("DataStructures")
local Dictionary, Stack, WordInfo, Environment = ds.Dictionary, ds.Stack, ds.WordInfo, ds.Environment
local helpers = require("CoreHelpers")
local areNumbers, getActiveDataStack, isNumber, isWhitespace= helpers.areNumbers, helpers.getActiveDataStack, helpers.isNumber, helpers.isWhitespace
local skipWhitespace, parseToken = helpers.skipWhitespace, helpers.parseToken
local popTwoOperands = helpers.popTwoOperands
local defineWord = helpers.defineWord
local ret = helpers.ret
local makeCall = helpers.makeCall

local literal = helpers.literal
local compile = helpers.compile
-- Mathematical operations
local function add(env: Environment)
  local a, b = popTwoOperands(env)
  if areNumbers(a,b) then
    local c = (a as number) + (b as number)
    getActiveDataStack(env):push(c)
  else
    error("invalid operands for add operation!")
  end
end
local function sub(env: Environment)
  local a, b = popTwoOperands(env)
  if areNumbers(a, b) then
    local c = (a as number) - (b as number)
    getActiveDataStack(env):push(c)
  else
    error("invalid operands for sub operation!")
  end
end
local function mul(env: Environment)
  local a, b = popTwoOperands(env)
  if areNumbers(a, b) then
    local c = (a as number) * (b as number)
    getActiveDataStack(env):push(c)
  else
    error("invalid operands for mul operation!")
  end
end
local function div(env: Environment)
  local a, b = popTwoOperands(env)
  if areNumbers(a, b) then
    local c = (a as number) / (b as number)
    getActiveDataStack(env):push(c)
  else
    error("invalid operands for div operation!")
  end
end

local function dup(env: Environment)
  local stack = getActiveDataStack(env)
  local top = stack:pop()
  stack:push(top)
  stack:push(top)
end
local function swap(env: Environment)
  local stack = getActiveDataStack(env)
  local a, b = popTwoOperands(env)
  stack:push(b)
  stack:push(a)
end
local function rot(env: Environment)
  local stack = getActiveDataStack(env)
  local c, b, a= stack:pop(), stack:pop(), stack:pop()
  stack:push(b)
  stack:push(c)
  stack:push(a)
end
local function drop(env: Environment)
  local stack = getActiveDataStack(env)
  stack:pop()
end
local function over(env: Environment)
  local stack = getActiveDataStack(env)
  local b, a = stack:pop(), stack:pop()
  stack:push(a)
  stack:push(b)
  stack:push(a)
end
-- I/O operations
local function dot(env: Environment)
	local out = env.activeDataStack:pop()
  io.write(tostring(out) as string)
end

local twoDup = helpers.makeCall{over, over, ret}

local function twoSwap(env: Environment)
  local stack = getActiveDataStack(env)
  local d, c, b, a = stack:pop(), stack:pop(), stack:pop(), stack:pop()
  stack:push(c)
  stack:push(d)
  stack:push(a)
  stack:push(b)
end


local function twoOver(env: Environment)
  local stack = getActiveDataStack(env)
  local d, c, b, a = stack:pop(), stack:pop(), stack:pop(), stack:pop()
  stack:push(a)
  stack:push(b)
  stack:push(c)
  stack:push(d)
  stack:push(a)
  stack:push(b)
end
local function nip(env: Environment)
  local stack = getActiveDataStack(env)
  local b, _ = stack:pop(), stack:pop()
  stack:push(b)
end

local tuck = helpers.makeCall{swap, over, ret}

local function roll(env: Environment)
  local stack = getActiveDataStack(env)
  local u = stack:pop()
  local bufferStack = Stack:new()
  if u is number then
    local v = u as number
    while(v > 0) do
      local item = stack:pop()
      bufferStack:push(item)
      v = v - 1
    end
    local newTop = stack:pop()
    local x = (u as number)
    while(x > 0) do
      local item = bufferStack:pop()
      stack:push(item)
      x = x - 1
    end
    stack:push(newTop)
  else
    error("u is not a number")
  end
end 
local function getExecutionToken(env: Environment)
  local stack = getActiveDataStack(env)
  skipWhitespace(env)
  local name: string = parseToken(env)
  for _, dictionary in ipairs(env.dictionaries) do
    local wordinfo = dictionary:lookup(name)
    if wordinfo then
      stack:push((wordinfo as WordInfo).func)
      break
    end
  end
end

local function execute(env: Environment)
  local stack = getActiveDataStack(env)
  local func: function(Environment) = stack:pop() as function(Environment)
  func(env)
end

local function enterCompileMode(e: Environment)
  e.compileState = true
end

local function exitCompileMode(e: Environment)
  e.compileState = false
end

local function colon(e: Environment)
  skipWhitespace(e)
  e.locals = Dictionary:new()
  local name: string | nil = parseToken(e)
  if not name then
    e.running = false
    return
  end
  e.currentDefinitionName = name as string
  enterCompileMode(e)
end

local function semicolon(e: Environment)
  local localCount = e.locals.wordCount
  local cleanup_locals = function(e: Environment)
    for _ = localCount, 1, -1 do
      e.returnStack:pop()
    end
  end
  table.insert(e.currentDefinition as {function(Environment)}, cleanup_locals)
  table.insert(e.currentDefinition as {function(Environment)}, helpers.ret)
  local instrs = e.currentDefinition

  local call = helpers.makeCall(instrs as {function(Environment)}, e.locals.wordCount)
  local dict = e.dictionaries[1]
  defineWord(dict, e.currentDefinitionName, call, false)
  e.mostRecentDefinition = e.currentDefinitionName
  e.currentDefinitionName = nil
  e.currentDefinition = {}
  -- clear locals
  e.locals = Dictionary:new()
  e.localBuffer = {}


  exitCompileMode(e)
end
local function toCompileStack(e: Environment)
  local dataStack = getActiveDataStack(e)
  local compileStack = e.compilerStack
  local val = dataStack:pop()
  compileStack:push(val)
end

local function fromCompileStack(e: Environment)
  e.activeDataStack:push(e.compilerStack:pop())
end

local function getDefinitionIndex(e: Environment)
  e.activeDataStack:push(#e.currentDefinition)
end

local function unresolvedZeroBranch(e: Environment)
  helpers.noop(e)
end

local function unresolvedUncondBranch(e: Environment)
  helpers.noop(e)
end

local function getResolvedZeroBranch(destinationIndex: integer): function(Environment)
  return function(e: Environment)
    local flag = e.activeDataStack:pop()
    if flag == 0 or not flag then
      e.instructionPointer.index = destinationIndex
    end
  end
end

local function resolveBranch(e: Environment, sourceIndex: integer, destinationIndex: integer): function(Environment)
  if e.currentDefinition[sourceIndex] == unresolvedZeroBranch then
    return function(e: Environment)
      local flag = e.activeDataStack:pop()
      if flag == 0 or not flag then
        e.instructionPointer.index = destinationIndex
      end
    end
  elseif e.currentDefinition[sourceIndex] == unresolvedUncondBranch then
    return function(e: Environment)
      e.instructionPointer.index = destinationIndex
    end
  end
  -- TODO: abort if none of the checks succeed
end


local function forthIf(e: Environment)
  table.insert(e.currentDefinition as {function(Environment)}, unresolvedZeroBranch)
  getDefinitionIndex(e)
  toCompileStack(e)
end

local function patchBranch(e: Environment)
  local branchIndex = e.activeDataStack:pop() as integer
  local destinationIndex = e.activeDataStack:pop() as integer
  e.currentDefinition[branchIndex as integer] = resolveBranch(e,branchIndex,destinationIndex)
end


local function forthThen(e: Environment)
  getDefinitionIndex(e)
  fromCompileStack(e)
  patchBranch(e)
end
local function forthElse(e: Environment)
  table.insert(e.currentDefinition as {function(Environment)}, unresolvedUncondBranch)
  getDefinitionIndex(e)
  fromCompileStack(e)
  patchBranch(e)
  getDefinitionIndex(e)
  toCompileStack(e)
end


local function cmplt(e: Environment)
  local a, b = popTwoOperands(e)
  e.activeDataStack:push((a as number) < (b as number))
end
local function cmplte(e: Environment)
  local a, b = popTwoOperands(e)
  e.activeDataStack:push((a as number) <= (b as number))
end
local function cmpe(e: Environment)
  local a, b = popTwoOperands(e)
  e.activeDataStack:push((a as number) == (b as number))
end
local function cmpgte(e: Environment)
  local a, b = popTwoOperands(e)
  e.activeDataStack:push((a as number) >= (b as number))
end
local function cmpgt(e: Environment)
  local a, b = popTwoOperands(e)
  e.activeDataStack:push((a as number) > (b as number))
end
local function cmpne(e: Environment)
  local a, b = popTwoOperands(e)
  e.activeDataStack:push((a as number) ~= (b as number))
end

local function markImmediate(e: Environment)
  helpers.searchDictionaries(e, e.mostRecentDefinition).immediate = true
end

local function forthBegin(e: Environment)
  getDefinitionIndex(e)
  toCompileStack(e)
end

local function forthUntil(e: Environment)
  fromCompileStack(e)
  table.insert(e.currentDefinition as {function(Environment)}, unresolvedZeroBranch)
  getDefinitionIndex(e)
  patchBranch(e)
end
local function forthAgain(e: Environment)
  fromCompileStack(e)
  table.insert(e.currentDefinition as {function(Environment)}, unresolvedUncondBranch)
  getDefinitionIndex(e)
  patchBranch(e)
end
local function forthWhile(e: Environment)
  fromCompileStack(e)
  table.insert(e.currentDefinition as {function(Environment)}, unresolvedZeroBranch)
  getDefinitionIndex(e)
  toCompileStack(e)
  toCompileStack(e)
end

local function forthRepeat(e: Environment)
  fromCompileStack(e)
  table.insert(e.currentDefinition as {function(Environment)}, unresolvedUncondBranch)
  getDefinitionIndex(e)
  patchBranch(e)
  getDefinitionIndex(e)
  fromCompileStack(e)
  patchBranch(e)
end

local function stringLiteral(e: Environment)
  local str: string | nil = ""
  local chr: string | nil = ""
  if not e.activeInputStream:curr() then e.activeDataStack:push("") end
  while(not (chr == "\"") and chr ~= nil) do
    str = str..(chr as string)
    chr = e.activeInputStream:next()
  end
  e.activeInputStream:next()
  e.activeDataStack:push(str as string)
end

local function forthTo(e: Environment)
  e.toMode = true
end

local function forthValue(e: Environment)
  local val = e.activeDataStack:pop() as any
  helpers.skipSpaces(e)
  local valName = helpers.parseToken(e)
  local index = #e.values + 1
  local getValue = function(e: Environment)
    e.activeDataStack:push(e.values[index])
  end
  local setValue = function(e: Environment)
    local newValue = e.activeDataStack:pop()
    e.values[index] = newValue
  end
  defineWord(e.dictionaries[1], valName, getValue, false, setValue)
  e.values[index] = val
end



local function forth_local_(e: Environment)
  local localName = e.activeDataStack:pop() as string
  if localName == "" then
    -- TODO: figure out what to do when you receive an "end of locals"
    local frame = #e.localBuffer
    for _, f in ipairs(e.localBuffer) do
      f(frame, e)
    end

  else
    -- local offset = e.locals.wordCount


    local index = #e.localBuffer + 1

    local postponedLocal = function(frame: integer, e: Environment)
      local localAddress = e.returnStack.top - frame + index
      local localFunction = function(e: Environment)
        local top = e.returnStack.top
        e.activeDataStack:push(e.returnStack.contents[top - frame + index])
      end
      local localFunctionTo = function(e: Environment)
        local top = e.returnStack.top
        local newValue = e.activeDataStack:pop()
        e.returnStack.contents[top - frame + index] = newValue
      end
      defineWord(e.locals, localName, localFunction, false, localFunctionTo)
    end
    table.insert(e.localBuffer, postponedLocal)
  end
end

local function parseName(e: Environment)
  helpers.skipSpaces(e)
  e.activeDataStack:push(helpers.parseToken(e))
end

local function emptys(e: Environment)
  e.activeDataStack:push("")
end

local CoreWords = Dictionary:new()
local addInfo = WordInfo:new(add, false)
local subInfo = WordInfo:new(sub, false)
local mulInfo = WordInfo:new(mul, false)
local divInfo = WordInfo:new(div, false)
local dotInfo = WordInfo:new(dot, false)
local dupInfo = WordInfo:new(dup, false)
local swapInfo = WordInfo:new(swap, false)
local rotInfo = WordInfo:new(rot, false)
local dropInfo = WordInfo:new(drop, false)
local overInfo = WordInfo:new(over, false)
local twoDupInfo = WordInfo:new(twoDup, false)
local twoSwapInfo = WordInfo:new(twoSwap, false)
local twoOverInfo = WordInfo:new(twoOver, false)
local nipInfo = WordInfo:new(nip, false)

defineWord(CoreWords, "TUCK", tuck, false)
defineWord(CoreWords, "ROLL", roll, false)
defineWord(CoreWords, "'", getExecutionToken, false)
defineWord(CoreWords, "EXECUTE", execute, false)
defineWord(CoreWords, "[", exitCompileMode, true)
defineWord(CoreWords, "]", enterCompileMode, false)
defineWord(CoreWords, "COMPILE,", compile, true)
defineWord(CoreWords, "LITERAL", literal, true)
defineWord(CoreWords, ":", colon, false)
defineWord(CoreWords, ";", semicolon, true)
defineWord(CoreWords, ">C",toCompileStack, false)
defineWord(CoreWords, "C>", fromCompileStack, false)
defineWord(CoreWords, "DP", getDefinitionIndex, false)
defineWord(CoreWords, "IF", forthIf, true)
defineWord(CoreWords, "THEN", forthThen, true)
defineWord(CoreWords, "ELSE", forthElse, true)
defineWord(CoreWords, "IMMEDIATE", markImmediate, false)
defineWord(CoreWords, "BEGIN", forthBegin, true)
defineWord(CoreWords, "UNTIL", forthUntil, true)
defineWord(CoreWords, "WHILE", forthWhile, true)
defineWord(CoreWords, "REPEAT", forthRepeat, true)
defineWord(CoreWords, "AGAIN", forthAgain, true)

defineWord(CoreWords, "TO", forthTo, true)
defineWord(CoreWords, "(LOCAL)", forth_local_, true)
defineWord(CoreWords, "PARSE-NAME", parseName, true)
defineWord(CoreWords, "EMPTYS", emptys, true)
defineWord(CoreWords, "\"\"", emptys, false)

defineWord(CoreWords, "VALUE", forthValue, false)

defineWord(CoreWords,"<", cmplt, false)
defineWord(CoreWords,"<=", cmplte, false)
defineWord(CoreWords,"=", cmpe, false)
defineWord(CoreWords,">=", cmpgte, false)
defineWord(CoreWords,">", cmpgt, false)
defineWord(CoreWords,"<>", cmpne, false)
defineWord(CoreWords,"S\"", stringLiteral, true)


CoreWords:define("+", addInfo)
CoreWords:define("-", subInfo)
CoreWords:define("*", mulInfo)
CoreWords:define("/", divInfo)
CoreWords:define(".", dotInfo)

CoreWords:define("DUP", dupInfo)
CoreWords:define("SWAP", swapInfo)
CoreWords:define("ROT", rotInfo)
CoreWords:define("DROP", dropInfo)
CoreWords:define("OVER", overInfo)
CoreWords:define("2DUP", twoDupInfo)
CoreWords:define("2SWAP", twoSwapInfo)
CoreWords:define("2OVER", twoOverInfo)
CoreWords:define("NIP", nipInfo)



local sqr = {dup, mul, ret}
local sqrf = makeCall(sqr)

defineWord(CoreWords, "SQR", sqrf, false)








return CoreWords