-- Client Startup Script for a Distributed Lua Shell Cluster
-- This script registers with the master and then listens for commands.
-- When a command is received (other than "cancel"), it executes it as Lua code.
-- The print function is overridden to capture output so that, for example,
-- print(9*9) returns "81" rather than a default return value.
-- IMPORTANT: Set the master computer's ID to 36.

local masterID = 36

-- Function to open all attached modems.
local function openAllModems()
  local sides = {"left", "right", "top", "bottom", "front", "back"}
  local opened = {}
  for _, side in ipairs(sides) do
    if peripheral.getType(side) == "modem" then
      rednet.open(side)
      table.insert(opened, side)
      print("Opened rednet on modem at side: " .. side)
    end
  end
  if #opened == 0 then
    error("No modems found! Please attach at least one modem.")
  end
end

openAllModems()

print("Client Computer ID is: " .. os.getComputerID())
print("Attempting to register with master (" .. masterID .. ") on protocol 'supercomp'...")

rednet.send(masterID, "register", "supercomp")
local senderID, message, protocol = rednet.receive("supercomp", 10)
if message == "registered" then
  print("Successfully registered with master (" .. masterID .. ").")
else
  print("Failed to register with master. Received: " .. tostring(message))
end

while true do
  local senderID, command, protocol = rednet.receive("supercomp")
  if command == "cancel" then
    -- Ignore cancel messages.
  else
    print("Received command: " .. tostring(command))
    local result, finished = nil, false

    local function executeCommand()
      -- Capture printed output in a table.
      local capturedOutput = {}
      local function capturePrint(...)
        local args = {...}
        local parts = {}
        for i, v in ipairs(args) do
          table.insert(parts, tostring(v))
        end
        table.insert(capturedOutput, table.concat(parts, "\t"))
      end

      -- Create a custom environment that uses our capturePrint for print.
      local env = {}
      for k, v in pairs(_G) do
        env[k] = v
      end
      env.print = capturePrint

      -- Try compiling the command as an expression first, then as a statement.
      local func, err = load("return " .. command, "command", "t", env)
      if not func then
        func, err = load(command, "command", "t", env)
      end

      if func then
        local success, res = pcall(func)
        if success then
          if #capturedOutput > 0 then
            result = table.concat(capturedOutput, "\n")
          else
            result = (res ~= nil and res or "nil")
          end
        else
          result = "Error: " .. tostring(res)
        end
      else
        result = "Error: " .. tostring(err)
      end
      finished = true
    end

    local function waitForCancel()
      local s, msg, p = rednet.receive("supercomp", 10)
      if msg == "cancel" then
        finished = false
        result = nil
      end
    end

    -- Run the command execution and cancellation listener concurrently.
    parallel.waitForAny(executeCommand, waitForCancel)
    if finished and result then
      rednet.send(masterID, tostring(result), "supercomp")
    end
  end
end