local parentShell = _ENV.shell

_ENV.shell = { }
_ENV.multishell = _ENV.multishell or { }

local fs         = _G.fs
local shell      = _ENV.shell
local multishell = _ENV.multishell

local sandboxEnv = setmetatable({ }, { __index = _G })
for k,v in pairs(_ENV) do
  sandboxEnv[k] = v
end
sandboxEnv.shell = shell
sandboxEnv.multishell = multishell

_G.requireInjector()

local Util = require('util')

local DIR = (parentShell and parentShell.dir()) or ""
local PATH = (parentShell and parentShell.path()) or ".:/rom/programs"
local ALIASES = (parentShell and parentShell.aliases()) or {}
local tCompletionInfo = (parentShell and parentShell.getCompletionInfo()) or {}

local bExit = false
local tProgramStack = {}

local function parseCommandLine( ... )
  local sLine = table.concat( { ... }, " " )
  local tWords = {}
  local bQuoted = false
  for match in string.gmatch( sLine .. "\"", "(.-)\"" ) do
    if bQuoted then
      table.insert( tWords, match )
    else
      for m in string.gmatch( match, "[^ \t]+" ) do
        table.insert( tWords, m )
      end
    end
    bQuoted = not bQuoted
  end

  return table.remove(tWords, 1), tWords
end

local function run(env, ...)
  local path, args = parseCommandLine(...)

  if not path then
    error('No such program')
  end

  local isUrl = not not path:match("^(https?:)//(([^/:]+):?([0-9]*))(/?.*)$")
  if not isUrl then
    path = shell.resolveProgram(path)
    if not path then
      error('No such program')
    end
  end

  local fn, err

  if isUrl then
    fn, err = Util.loadUrl(path, env)
  else
    fn, err = loadfile(path, env)
  end

  if not fn then
    error(err)
  end

  local oldTitle

  if multishell and multishell.getTitle then
    oldTitle = multishell.getTitle(multishell.getCurrent())
    multishell.setTitle(multishell.getCurrent(), fs.getName(path))
  end

  if isUrl then
    tProgramStack[#tProgramStack + 1] = path:match("^https?://([^/:]+:?[0-9]*/?.*)$")
  else
    tProgramStack[#tProgramStack + 1] = path
  end
  local r = { fn(table.unpack(args)) }

  tProgramStack[#tProgramStack] = nil

  if multishell and multishell.getTitle then
    multishell.setTitle(multishell.getCurrent(), oldTitle or 'shell')
  end

  return table.unpack(r)
end

-- Install shell API
function shell.run(...)
  return pcall(run, setmetatable(Util.shallowCopy(sandboxEnv), { __index = _G }), ...)
end

function shell.exit()
  bExit = true
end

function shell.dir() return DIR end
function shell.setDir(d) DIR = d end
function shell.path() return PATH end
function shell.setPath(p) PATH = p end

function shell.resolve( _sPath )
  local sStartChar = string.sub( _sPath, 1, 1 )
  if sStartChar == "/" or sStartChar == "\\" then
    return fs.combine( "", _sPath )
  else
    return fs.combine(DIR, _sPath )
  end
end

function shell.resolveProgram( _sCommand )

  if ALIASES[ _sCommand ] ~= nil then
    _sCommand = ALIASES[ _sCommand ]
  end

  local path = shell.resolve(_sCommand)
  if fs.exists(path) and not fs.isDir(path) then
    return path
  end
  if fs.exists(path .. '.lua') then
    return path .. '.lua'
  end

  -- If the path is a global path, use it directly
  local sStartChar = string.sub( _sCommand, 1, 1 )
  if sStartChar == "/" or sStartChar == "\\" then
    local sPath = fs.combine( "", _sCommand )
    if fs.exists( sPath ) and not fs.isDir( sPath ) then
      return sPath
    end
    return nil
  end

  -- Otherwise, look on the path variable
  for sPath in string.gmatch(PATH or '', "[^:]+") do
    sPath = fs.combine(sPath, _sCommand )
    if fs.exists( sPath ) and not fs.isDir( sPath ) then
      return sPath
    end
    if fs.exists(sPath .. '.lua') then
      return sPath .. '.lua'
    end
  end
  -- Not found
  return nil
end

function shell.programs( _bIncludeHidden )
  local tItems = {}

  -- Add programs from the path
  for sPath in string.gmatch(PATH, "[^:]+") do
    sPath = shell.resolve(sPath)
    if fs.isDir( sPath ) then
      local tList = fs.list( sPath )
      for _,sFile in pairs( tList ) do
        if not fs.isDir( fs.combine( sPath, sFile ) ) and
          (_bIncludeHidden or string.sub( sFile, 1, 1 ) ~= ".") then
          tItems[ sFile ] = true
        end
      end
    end
  end

  -- Sort and return
  local tItemList = {}
  for sItem in pairs( tItems ) do
    table.insert( tItemList, sItem )
  end
  table.sort( tItemList )
  return tItemList
end

function shell.complete(sLine) end
function shell.completeProgram(sProgram) end

function shell.setCompletionFunction(sProgram, fnComplete)
  tCompletionInfo[sProgram] = { fnComplete = fnComplete }
end

function shell.getCompletionInfo()
  return tCompletionInfo
end

function shell.getRunningProgram()
  return tProgramStack[#tProgramStack]
end

function shell.setAlias( _sCommand, _sProgram )
  ALIASES[ _sCommand ] = _sProgram
end

function shell.clearAlias( _sCommand )
  ALIASES[ _sCommand ] = nil
end

function shell.aliases()
  local tCopy = {}
  for sAlias, sCommand in pairs(ALIASES) do
    tCopy[sAlias] = sCommand
  end
  return tCopy
end

function shell.newTab(tabInfo, ...)
  local path, args = parseCommandLine(...)
  path = shell.resolveProgram(path)

  if path then
    tabInfo.path = path
    tabInfo.env = sandboxEnv
    tabInfo.args = Util.shallowCopy(args)
    tabInfo.title = fs.getName(path)

    if path ~= 'sys/apps/shell' then
      table.insert(tabInfo.args, 1, tabInfo.path)
      tabInfo.path = 'sys/apps/shell'
    end
    return multishell.openTab(tabInfo)
  end
  return nil, 'No such program'
end

function shell.openTab( ... )
  return shell.newTab({ }, ...)
end

function shell.openForegroundTab( ... )
  return shell.newTab({ focused = true }, ...)
end

function shell.openHiddenTab( ... )
  return shell.newTab({ hidden = true }, ...)
end

function shell.switchTab(tabId)
  multishell.setFocus(tabId)
end

local tArgs = { ... }
if #tArgs > 0 then
  return run(_ENV, ...)
end

local Config  = require('config')
local History = require('history')

local colors    = _G.colors
local keys      = _G.keys
local os        = _G.os
local term      = _G.term
local textutils = _G.textutils

local config = {
  standard = {
    textColor  = colors.white,
    commandTextColor = colors.lightGray,
    directoryTextColor = colors.gray,
    directoryBackgroundColor = colors.black,
    promptTextColor = colors.gray,
    promptBackgroundColor = colors.black,
    directoryColor = colors.gray,
  },
  color = {
    textColor = colors.white,
    commandTextColor = colors.yellow,
    directoryTextColor  = colors.orange,
    directoryBackgroundColor = colors.black,
    promptTextColor = colors.blue,
    promptBackgroundColor = colors.black,
    directoryColor = colors.green,
  },
  displayDirectory = true,
}

Config.load('shellprompt', config)

local _colors = config.standard
if term.isColor() then
  _colors = config.color
end

local function autocompleteFile(results, words)

  local function getBaseDir(path)
    if #path > 1 then
      if path:sub(-1) ~= '/' then
        path = fs.getDir(path)
      end
    end
    if path:sub(1, 1) == '/' then
      path = fs.combine(path, '')
    else
      path = fs.combine(shell.dir(), path)
    end
    while not fs.isDir(path) do
      path = fs.getDir(path)
    end
    return path
  end

  local function getRawPath(path)
    local baseDir = ''
    if path:sub(1, 1) ~= '/' then
      baseDir = shell.dir()
    end
    if #path > 1 then
      if path:sub(-1) ~= '/' then
        path = fs.getDir(path)
      end
    end
    if fs.isDir(fs.combine(baseDir, path)) then
      return path
    end
    return fs.getDir(path)
  end

  local match = words[#words] or ''
  local startDir = getBaseDir(match)
  local rawPath = getRawPath(match)

  if fs.isDir(startDir) then
    local files = fs.list(startDir)
    for _,f in pairs(files) do
      local path = fs.combine(rawPath, f)
      if fs.isDir(fs.combine(startDir, f)) then
        results[path .. '/'] = 'directory'
      else
        results[path .. ' '] = 'program'
      end
    end
  end
end

local function autocompleteProgram(results, words)
  if #words == 1 then
    local files = shell.programs(true)
    for _,f in ipairs(files) do
      results[f .. ' '] = 'program'
    end
    for f in pairs(ALIASES) do
      results[f .. ' '] = 'program'
    end
  end
end

local function autocompleteArgument(results, program, words)
  local word = ''
  if #words > 1 then
    word = words[#words]
  end

  local tInfo = tCompletionInfo[program]
  local args = tInfo.fnComplete(shell, #words - 1, word, words)
  if args then
    Util.filterInplace(args, function(f)
      return not Util.key(args, f .. '/')
    end)
    for _,arg in ipairs(args) do
      results[word .. arg] = 'argument'
    end
  end
end

local function autocomplete(line, suggestions)
  local words = { }
  for word in line:gmatch("%S+") do
    table.insert(words, word)
  end
  if line:match(' $') then
    table.insert(words, '')
  end

  local results = { }
  local files = { }

  if #words == 0 then
    files = autocompleteFile(results, words)
  else
    local program = shell.resolveProgram(words[1])
    if tCompletionInfo[program] then
      autocompleteArgument(results, program, words)
    else
      autocompleteProgram(results, words)
      autocompleteFile(results, words)
    end
  end

  local match = words[#words] or ''
  for f in pairs(results) do
    if f:sub(1, #match) == match then
      table.insert(files, f)
    end
  end

  if #files == 1 then
    words[#words] = files[1]
    return table.concat(words, ' ')
  elseif #files > 1 and suggestions then
    print()

    local word = words[#words] or ''
    local prefix = word:match("(.*/)") or ''
    if #prefix > 0 then
      for _,f in ipairs(files) do
        if f:match("^" .. prefix) ~= prefix then
          prefix = ''
          break
        end
      end
    end

    local tDirs, tFiles = { }, { }
    for _,f in ipairs(files) do
      if results[f] == 'directory' then
        f = f:gsub(prefix, '', 1)
        table.insert(tDirs, f)
      else
        f = f:gsub(prefix, '', 1)
        table.insert(tFiles, f)
      end
    end
    table.sort(tDirs)
    table.sort(tFiles)

    if #tDirs > 0 and #tDirs < #tFiles then
      local w = term.getSize()
      local nMaxLen = w / 8
      for _,sItem in pairs(files) do
        nMaxLen = math.max(string.len(sItem) + 1, nMaxLen)
      end
      local nCols = math.floor(w / nMaxLen)
      if #tDirs < nCols then
        for _ = #tDirs + 1, nCols do
          table.insert(tDirs, '')
        end
      end
    end

    if #tDirs > 0 then
      textutils.tabulate(_colors.directoryColor, tDirs, colors.white, tFiles)
    else
      textutils.tabulate(colors.white, tFiles)
    end

    term.setTextColour(_colors.promptTextColor)
    term.setBackgroundColor(_colors.promptBackgroundColor)
    term.write("$ " )

    term.setTextColour(_colors.commandTextColor)
    term.setBackgroundColor(colors.black)
    return line
  elseif #files > 1 then

    -- ugly (complete as much as possible)
    local word = words[#words] or ''
    local i = #word + 1
    while true do
      local ch
      for _,f in ipairs(files) do
        if #f < i then
          words[#words] = string.sub(f, 1, i - 1)
          return table.concat(words, ' ')
        end
        if not ch then
          ch = string.sub(f, i, i)
        elseif string.sub(f, i, i) ~= ch then
          if i == #word + 1 then
            return
          end
          words[#words] = string.sub(f, 1, i - 1)
          return table.concat(words, ' ')
        end
      end
      i = i + 1
    end
  end
end

local function shellRead(history)
  term.setCursorBlink( true )

  local sLine = ""
  local nPos = 0
  local lastPattern

  local w = term.getSize()
  local sx = term.getCursorPos()

  history:reset()

  local function redraw( sReplace )
    local nScroll = 0
    if sx + nPos >= w then
      nScroll = (sx + nPos) - w
    end

    local _,cy = term.getCursorPos()
    term.setCursorPos( sx, cy )
    if sReplace then
      term.write( string.rep( sReplace, math.max( string.len(sLine) - nScroll, 0 ) ) )
    else
      term.write( string.sub( sLine, nScroll + 1 ) )
    end
    term.setCursorPos( sx + nPos - nScroll, cy )
  end

  while true do
    local sEvent, param = os.pullEventRaw()

    if sEvent == "char" then
      sLine = string.sub( sLine, 1, nPos ) .. param .. string.sub( sLine, nPos + 1 )
      nPos = nPos + 1
      redraw()
    elseif sEvent == "paste" then
      sLine = string.sub( sLine, 1, nPos ) .. param .. string.sub( sLine, nPos + 1 )
      nPos = nPos + string.len( param )
      redraw()
    elseif sEvent == 'mouse_click' and param == 2 then
      redraw(string.rep(' ', #sLine))
      sLine = ''
      nPos = 0
      redraw()
    elseif sEvent == 'terminate' then
      bExit = true
      break
    elseif sEvent == "key" then
      if param == keys.enter then
        -- Enter
        break
      elseif param == keys.tab then
        if nPos == #sLine then
          local showSuggestions = lastPattern == sLine
          lastPattern = sLine

          local cline = autocomplete(sLine, showSuggestions)
          if cline then
            sLine = cline
            nPos = #sLine
            redraw()
          end
        end
      elseif param == keys.left then
        if nPos > 0 then
          nPos = nPos - 1
          redraw()
        end
      elseif param == keys.right then
        if nPos < string.len(sLine) then
          redraw(" ")
          nPos = nPos + 1
          redraw()
        end
      elseif param == keys.up or param == keys.down then
        redraw(" ")
        if param == keys.up then
          sLine = history:back()
        else
          sLine = history:forward()
        end
        if sLine then
          nPos = string.len(sLine)
        else
          sLine = ""
          nPos = 0
        end
        redraw()
      elseif param == keys.backspace then
        if nPos > 0 then
          redraw(" ")
          sLine = string.sub( sLine, 1, nPos - 1 ) .. string.sub( sLine, nPos + 1 )
          nPos = nPos - 1
          redraw()
        end
      elseif param == keys.home then
        redraw(" ")
        nPos = 0
        redraw()
      elseif param == keys.delete then
        if nPos < string.len(sLine) then
          redraw(" ")
          sLine = string.sub( sLine, 1, nPos ) .. string.sub( sLine, nPos + 2 )
          redraw()
        end
      elseif param == keys["end"] then
        redraw(" ")
        nPos = string.len(sLine)
        redraw()
      end
    elseif sEvent == "term_resize" then
      w = term.getSize()
      redraw()
    end
  end

  local _, cy = term.getCursorPos()
  term.setCursorPos( w + 1, cy )
  print()
  term.setCursorBlink( false )
  return sLine
end

local history = History.load('usr/.shell_history', 25)

while not bExit do
  if config.displayDirectory then
    term.setTextColour(_colors.directoryTextColor)
    term.setBackgroundColor(_colors.directoryBackgroundColor)
    print('==' .. os.getComputerLabel() .. ':/' .. DIR)
  end
  term.setTextColour(_colors.promptTextColor)
  term.setBackgroundColor(_colors.promptBackgroundColor)
  term.write("$ " )
  term.setTextColour(_colors.commandTextColor)
  term.setBackgroundColor(colors.black)
  local sLine = shellRead(history)
  if bExit then -- terminated
    break
  end
  sLine = Util.trim(sLine)
  if #sLine > 0 and sLine ~= 'exit' then
    history:add(sLine)
  end
  term.setTextColour(_colors.textColor)
  if #sLine > 0 then
    local result, err = shell.run(sLine)
    if not result and err then
      _G.printError(err)
    end
  end
end