-- Computer-side MCP execution primitives.
--
-- Protocol-agnostic building blocks for the MCP `exec-lua` / `write-file` features.
-- The transport lives elsewhere: servers/mcp-computer-server.lua reacts to sandbox
-- gateway_request events and replies through the trap-sandbox gateway (see
-- apis/libsandbox.lua). This module only knows how to run code and write files.

local function defaultOs()
  return os;
end

local function formatLabel(label)
  if type(label) ~= 'string' or label == '' then
    return 'null';
  end
  return label;
end

local function appendOutput(output, value)
  output[#output + 1] = tostring(value);
end

local function normalizeArgs(rawArgs)
  if rawArgs == nil then
    return true, {};
  end
  if type(rawArgs) ~= 'table' then
    return false, 'args must be a table of strings';
  end

  local args = {};
  for i = 1, #rawArgs do
    local arg = rawArgs[i];
    if type(arg) ~= 'string' then
      return false, 'args must be a table of strings';
    end
    args[i] = arg;
  end
  return true, args;
end

local function copyArgs(args)
  local copy = {};
  for i = 1, #args do
    copy[i] = args[i];
  end
  return copy;
end

local function dirname(path)
  local dir = string.match(path, '^(.*)/[^/]*$');
  if dir == nil then
    return '';
  end
  if dir == '' then
    return '/';
  end
  return dir;
end

local function createShellProxy(shellLib, runningProgram)
  if type(shellLib) ~= 'table' then
    return nil;
  end

  return setmetatable({
    getRunningProgram = function()
      return runningProgram;
    end,
  }, { __index = shellLib });
end

local function serializeReturn(value)
  local valueType = type(value);
  if valueType == 'nil' then
    return { type = 'nil' };
  end

  if valueType == 'string' or valueType == 'number' or valueType == 'boolean' then
    return { type = valueType, value = value };
  end

  local ok, serialized = pcall(textutils.serialize, value);
  if ok then
    return { type = valueType, repr = serialized };
  end
  return { type = valueType, repr = tostring(value) };
end

-- Builds a CraftOS-style require/package pair the same way /rom/programs/shell.lua does
-- (dofile of cc.require, then .make(env, dir)). Returned factory caches the make function so
-- the rom module is only loaded once. Stays synchronous (require never yields), so it is safe
-- inside the synchronous sandbox handler -- unlike spawning a shell, which yields and deadlocks.
local function defaultRequireFactory()
  local make;
  local loaded = false;
  return function(env, dir)
    if not loaded then
      loaded = true;
      local ok, mod = pcall(dofile, '/rom/modules/main/cc/require.lua');
      if ok and type(mod) == 'table' then
        make = mod.make;
      end
    end
    if type(make) ~= 'function' then
      return nil, nil;
    end
    local madeOk, req, pkg = pcall(make, env, dir or '');
    if not madeOk then
      return nil, nil;
    end
    return req, pkg;
  end
end

local function createExecEnv(output, requireFactory, requireDir, execOpts)
  execOpts = execOpts or {};
  local env = setmetatable({}, { __index = _G });
  env.arg = copyArgs(execOpts.args or {});

  env.write = function(value)
    appendOutput(output, value);
  end;

  env.print = function(...)
    local values = table.pack(...);
    for i = 1, values.n do
      if i > 1 then
        appendOutput(output, '\t');
      end
      appendOutput(output, values[i]);
    end
    appendOutput(output, '\n');
  end;

  if requireFactory then
    local req, pkg = requireFactory(env, requireDir);
    if req then
      env.require = req;
      env.package = pkg;
    end
  end

  local shellProxy = createShellProxy(execOpts.shell, execOpts.runningProgram);
  if shellProxy then
    env.shell = shellProxy;
  end

  return env;
end

local function executeSource(code, chunkName, deps, execOpts)
  execOpts = execOpts or {};
  local output = {};
  local args = execOpts.args or {};

  local fn, loadErr = load(code, chunkName, 't', createExecEnv(output, deps.requireFactory, deps.requireDir, execOpts));
  if not fn then
    return { ok = false, error = tostring(loadErr), output = table.concat(output) };
  end

  local values = table.pack(pcall(fn, table.unpack(args)));
  if not values[1] then
    return { ok = false, error = tostring(values[2]), output = table.concat(output) };
  end

  local returns = {};
  for i = 2, values.n do
    returns[#returns + 1] = serializeReturn(values[i]);
  end
  return { ok = true, returns = returns, output = table.concat(output) };
end

local function createMcpComputer(deps)
  deps = deps or {};

  -- requireFactory: function(env, dir) -> require, package. Defaults to the rom-backed factory;
  -- pass `false` to disable require injection, or a function to stub it in tests.
  local requireFactory = deps.requireFactory;
  if requireFactory == nil then
    requireFactory = defaultRequireFactory();
  end
  local resolved = {
    requireFactory = requireFactory or nil,
    requireDir = deps.requireDir or '',
    shell = deps.shell or shell,
  };

  local api = {};

  api.formatLabel = formatLabel;

  function api.formatPong(osLike)
    osLike = osLike or defaultOs();
    return 'pong from ' .. tostring(osLike.getComputerID())
      .. ' (Label: ' .. formatLabel(osLike.getComputerLabel()) .. ')';
  end

  function api.executeLua(code)
    if type(code) ~= 'string' or code == '' then
      return { ok = false, error = 'code must be a non-empty string', output = '' };
    end

    return executeSource(code, 'mcp-exec', resolved);
  end

  function api.executeFile(path, fsLike, rawArgs)
    fsLike = fsLike or fs;
    if type(path) ~= 'string' or path == '' then
      return { ok = false, error = 'path must be a non-empty string', output = '' };
    end
    local argsOk, argsOrErr = normalizeArgs(rawArgs);
    if not argsOk then
      return { ok = false, error = argsOrErr, output = '' };
    end

    local handle, openErr = fsLike.open(path, 'r');
    if not handle then
      return { ok = false, error = tostring(openErr or 'failed to open file'), output = '' };
    end

    local ok, contentOrErr = pcall(function()
      return handle.readAll();
    end);
    local closeOk, closeErr = pcall(function()
      handle.close();
    end);

    if not ok then
      return { ok = false, error = tostring(contentOrErr), output = '' };
    end
    if not closeOk then
      return { ok = false, error = tostring(closeErr), output = '' };
    end

    local fileDeps = {
      requireFactory = resolved.requireFactory,
      requireDir = dirname(path),
    };
    return executeSource(contentOrErr, '@' .. path, fileDeps, {
      args = argsOrErr,
      runningProgram = path,
      shell = resolved.shell,
    });
  end

  function api.writeFile(path, content, fsLike)
    fsLike = fsLike or fs;
    if type(path) ~= 'string' or path == '' then
      return { ok = false, error = 'path must be a non-empty string' };
    end

    if type(content) ~= 'string' then
      return { ok = false, error = 'content must be a string' };
    end

    local handle, openErr = fsLike.open(path, 'w');
    if not handle then
      return { ok = false, error = tostring(openErr or 'failed to open file') };
    end

    local ok, writeErr = pcall(function()
      handle.write(content);
    end);
    local closeOk, closeErr = pcall(function()
      handle.close();
    end);

    if not ok then
      return { ok = false, error = tostring(writeErr) };
    end
    if not closeOk then
      return { ok = false, error = tostring(closeErr) };
    end

    return { ok = true, path = path, bytes = string.len(content) };
  end

  return api;
end

return createMcpComputer;