-- pushfs_server.lua
-- CC:Tweaked turtle/computer daemon to receive files and run remote commands.
-- Ops: ping | put | begin/chunk/end | exec
-- Protocol: "pushfs-v1"

--------------------------
-- Config
--------------------------
local PROTOCOL = "pushfs-v1"
local SHARED_SECRET = nil         -- set to a string to require a token
local MAX_IDLE_SEC = 120
local MAX_RETURN_BYTES = 64*1024  -- cap returned output

-- Lua 5.1/5.2 compat
local unpack = table.unpack or unpack

--------------------------
-- Utilities
--------------------------
local function openModems()
  local ok = false
  for _, s in ipairs(rs.getSides()) do
    if peripheral.getType(s) == "modem" then
      if not rednet.isOpen(s) then rednet.open(s) end
      ok = true
    end
  end
  if not ok then error("No modem found. Attach a wireless modem.") end
end

local function adler32(str)
  local MOD = 65521
  local a,b = 1,0
  for i=1,#str do
    a = (a + string.byte(str,i)) % MOD
    b = (b + a) % MOD
  end
  return string.format("%08x", bit32.bor(bit32.lshift(b,16), a))
end

local function readAll(path)
  if not fs.exists(path) or fs.isDir(path) then return "" end
  local h = fs.open(path, "rb"); local d = h.readAll() or ""; h.close(); return d
end

local function parentDir(path)
  local p = path:gsub("/+$","")
  local i = p:match("^(.*)/[^/]+$")
  return i or ""
end

local function sanitizePath(p)
  if p:sub(1,1) ~= "/" then p = "/"..p end
  p = p:gsub("//+","/")
  if p:find("%.%.") then return nil, "Path traversal not allowed" end
  if p:sub(1,5) == "/rom/" or p == "/rom" then return nil, "ROM is read-only" end
  return p
end

local function matchesTarget(msg, myId, myLabel)
  local t = msg.target
  if t == "all" then return true end
  if type(t) == "table" then
    if t.id and t.id == myId then return true end
    if t.label and myLabel and t.label == myLabel then return true end
  end
  return false
end

local function ack(toId, txid, ok, op, note, extra)
  local m = { proto=PROTOCOL, txid=txid, ok=ok, op=op, note=note }
  if extra then for k,v in pairs(extra) do m[k]=v end end
  rednet.send(toId, m, PROTOCOL)
end

--------------------------
-- File transfers
--------------------------
local xfers = {}  -- txid -> { tmp, dst, size, checksum, received, h, t_start }

local function cleanupTx(tx)
  if tx.h then pcall(function() tx.h.close() end) end
  if tx.tmp and fs.exists(tx.tmp) then pcall(function() fs.delete(tx.tmp) end) end
end

local function handle_begin(sender, msg)
  if SHARED_SECRET and msg.token ~= SHARED_SECRET then
    return ack(sender, msg.txid, false, "begin", "auth failed")
  end
  if type(msg.fname) ~= "string" or type(msg.size) ~= "number" then
    return ack(sender, msg.txid, false, "begin", "bad params")
  end
  local dst, err = sanitizePath(msg.fname)
  if not dst then return ack(sender, msg.txid, false, "begin", err) end

  local tmp = dst..".push.tmp"
  local pdir = parentDir(dst)
  if pdir ~= "" and not fs.exists(pdir) then fs.makeDir(pdir) end
  if fs.exists(tmp) then fs.delete(tmp) end

  local h, e = fs.open(tmp, "wb")
  if not h then return ack(sender, msg.txid, false, "begin", "open tmp failed: "..tostring(e)) end

  xfers[msg.txid] = {
    tmp = tmp, dst = dst, size = msg.size,
    checksum = msg.checksum, received = 0, h = h, t_start = os.clock()
  }
  return ack(sender, msg.txid, true, "begin", "ready")
end

local function handle_chunk(sender, msg)
  local tx = xfers[msg.txid]
  if not tx then return ack(sender, msg.txid, false, "chunk", "unknown txid") end
  if not msg.data or type(msg.data) ~= "string" then
    return ack(sender, msg.txid, false, "chunk", "no data")
  end
  tx.h.write(msg.data)
  tx.received = tx.received + #msg.data
  tx.t_start = os.clock()
  return ack(sender, msg.txid, true, "chunk", "ok", {received=tx.received})
end

local function handle_end(sender, msg)
  local tx = xfers[msg.txid]
  if not tx then return ack(sender, msg.txid, false, "end", "unknown txid") end

  tx.h.close(); tx.h = nil

  if tx.size and tx.received ~= tx.size then
    cleanupTx(tx); xfers[msg.txid]=nil
    return ack(sender, msg.txid, false, "end", ("size mismatch %d!=%d"):format(tx.received, tx.size))
  end

  if tx.checksum then
    local data = readAll(tx.tmp)
    local sum = adler32(data)
    if sum ~= tx.checksum then
      cleanupTx(tx); xfers[msg.txid]=nil
      return ack(sender, msg.txid, false, "end", "checksum mismatch", {have=sum, want=tx.checksum})
    end
  end

  if fs.exists(tx.dst) and fs.isDir(tx.dst) then
    cleanupTx(tx); xfers[msg.txid]=nil
    return ack(sender, msg.txid, false, "end", "destination is a directory")
  end
  if fs.exists(tx.dst) then fs.delete(tx.dst) end
  fs.move(tx.tmp, tx.dst)

  xfers[msg.txid]=nil
  return ack(sender, msg.txid, true, "end", "written", {path=tx.dst, bytes=tx.received})
end

local function handle_put(sender, msg)
  if SHARED_SECRET and msg.token ~= SHARED_SECRET then
    return ack(sender, msg.txid, false, "put", "auth failed")
  end
  if type(msg.fname) ~= "string" or type(msg.data) ~= "string" then
    return ack(sender, msg.txid, false, "put", "bad params")
  end
  local dst, err = sanitizePath(msg.fname)
  if not dst then return ack(sender, msg.txid, false, "put", err) end
  local pdir = parentDir(dst)
  if pdir ~= "" and not fs.exists(pdir) then fs.makeDir(pdir) end

  local tmp = dst..".push.tmp"
  if fs.exists(tmp) then fs.delete(tmp) end
  local h, e = fs.open(tmp, "wb")
  if not h then return ack(sender, msg.txid, false, "put", "open tmp failed: "..tostring(e)) end
  h.write(msg.data); h.close()

  if msg.checksum then
    local have = adler32(readAll(tmp))
    if have ~= msg.checksum then
      fs.delete(tmp)
      return ack(sender, msg.txid, false, "put", "checksum mismatch", {have=have, want=msg.checksum})
    end
  end

  if fs.exists(dst) then fs.delete(dst) end
  fs.move(tmp, dst)
  return ack(sender, msg.txid, true, "put", "written", {path=dst, bytes=#msg.data})
end

--------------------------
-- Exec
--------------------------
-- tiny in-memory terminal to capture output
local function makeBufferTerm(width, height)
  local buf = {}
  local x,y,tx,bc = 1,1,colors.white,colors.black
  local w,h = width or 80, height or 25
  for i=1,h do buf[i] = string.rep(" ", w) end

  local function clamp()
    if x < 1 then x = 1 elseif x > w then x = w end
    if y < 1 then y = 1 elseif y > h then y = h end
  end

  local t = {}
  function t.write(s)
    s = tostring(s)
    local line = buf[y]
    local left = line:sub(1, x-1)
    local right = line:sub(x + #s)
    buf[y] = left .. s .. right
    x = x + #s
    clamp()
  end
  function t.blit(s, tc, bc_) t.write(s) end
  function t.clear() for i=1,h do buf[i] = string.rep(" ", w) end x,y=1,1 end
  function t.clearLine() buf[y] = string.rep(" ", w); x=1 end
  function t.getCursorPos() return x,y end
  function t.setCursorPos(nx,ny) x,y=nx,ny; clamp() end
  function t.setCursorBlink(_) end
  function t.getSize() return w,h end
  function t.scroll(n)
    n = n or 1
    if n > 0 then for _=1,n do table.remove(buf,1); table.insert(buf, string.rep(" ", w)) end
    elseif n < 0 then for _=1,-n do table.remove(buf,#buf); table.insert(buf,1, string.rep(" ", w)) end
    end
  end
  function t.isColor() return term.isColor and term.isColor() or true end
  function t.setTextColor(c) tx=c end
  function t.setBackgroundColor(c) bc=c end
  function t.concat() return table.concat(buf, "\n") end
  return t
end

local function split_cmdline(line)
  local out, i, n = {}, 1, #line
  while i <= n do
    while i<=n and line:sub(i,i):match("%s") do i=i+1 end
    if i>n then break end
    local tok = {}
    local quote = nil
    while i<=n do
      local ch = line:sub(i,i)
      if quote then
        if ch == quote then quote=nil; i=i+1; break
        else tok[#tok+1]=ch; i=i+1 end
      else
        if ch == '"' or ch == "'" then quote = ch; i=i+1
        elseif ch:match("%s") then i=i+1; break
        else tok[#tok+1]=ch; i=i+1 end
      end
    end
    out[#out+1] = table.concat(tok)
  end
  return out
end

local function run_command_capture(line, cwd, timeout_s)
  local w,h = term.getSize()
  local cap = makeBufferTerm(w,h)
  local oldTerm = term.redirect(cap)
  local oldDir = shell.dir()
  if cwd and type(cwd)=="string" and fs.isDir(cwd) then shell.setDir(cwd) end

  local ok, err
  local function runner()
    local args = split_cmdline(line)
    if #args == 0 then return true end
    ok = shell.run(unpack(args))
    if not ok then err = "command failed" end
  end

  local done = false
  local function wrap()
    local s, e = pcall(runner)
    if not s then ok=false; err=tostring(e) end
    done = true
  end

  if timeout_s and timeout_s > 0 then
    local tid = os.startTimer(timeout_s)
    parallel.waitForAny(function() wrap() end, function()
      while true do local ev,id=os.pullEvent("timer"); if id==tid then return end end
    end)
    if not done then ok=false; err="timeout" end
  else
    wrap()
  end

  term.redirect(oldTerm)
  shell.setDir(oldDir)
  local out = cap.concat()
  if #out > MAX_RETURN_BYTES then
    out = out:sub(1, MAX_RETURN_BYTES) .. "\n[output truncated]"
  end
  return ok == true, out, err
end

local function handle_exec(sender, msg)
  -- msg: {op="exec", txid, line="<cmd>", capture=true/false, timeout=sec, cwd="/", token=...}
  if SHARED_SECRET and msg.token ~= SHARED_SECRET then
    return ack(sender, msg.txid, false, "exec", "auth failed")
  end
  if type(msg.line) ~= "string" or #msg.line == 0 then
    return ack(sender, msg.txid, false, "exec", "missing line")
  end

  if msg.capture then
    local ok, out, err = run_command_capture(msg.line, msg.cwd, msg.timeout)
    return ack(sender, msg.txid, ok, "exec", ok and "ok" or (err or "error"), {output=out})
  else
    -- background only when tabs exist; otherwise reject so client can use capture
    local cmdline = msg.line
    if multishell and multishell.launch then
      multishell.launch({}, "/rom/programs/shell", "-c", cmdline)
      return ack(sender, msg.txid, true, "exec", "started")
    elseif shell.openTab then
      shell.openTab("/rom/programs/shell", "-c", cmdline)
      return ack(sender, msg.txid, true, "exec", "started")
    else
      return ack(sender, msg.txid, false, "exec", "no background support on this device; use capture=true")
    end
  end
end

local function handle_ping(sender, msg, myId, myLabel)
  return ack(sender, msg.txid, true, "ping", "pong", {id=myId, label=myLabel or "", time=os.epoch("utc")})
end

--------------------------
-- Reaper
--------------------------
local function reapStale()
  local now = os.clock()
  for txid, tx in pairs(xfers) do
    if now - (tx.t_start or now) > MAX_IDLE_SEC then
      cleanupTx(tx)
      xfers[txid] = nil
    end
  end
end

--------------------------
-- Main
--------------------------
local function main()
  openModems()
  local myId = os.getComputerID()
  local myLabel = os.getComputerLabel()
  rednet.host(PROTOCOL, (myLabel or ("node-"..tostring(myId))))

  print(("pushfs_server: id=%d label=%s"):format(myId, tostring(myLabel)))
  print("Waiting...")

  while true do
    local ev = { os.pullEvent() }
    if ev[1] == "rednet_message" then
      local sender, message, proto = ev[2], ev[3], ev[4]
      if proto == PROTOCOL and type(message) == "table" and matchesTarget(message, myId, myLabel) then
        local op = message.op
        if     op == "begin" then handle_begin(sender, message)
        elseif op == "chunk" then handle_chunk(sender, message)
        elseif op == "end"   then handle_end(sender, message)
        elseif op == "put"   then handle_put(sender, message)
        elseif op == "exec"  then handle_exec(sender, message)
        elseif op == "ping"  then handle_ping(sender, message, myId, myLabel)
        else ack(sender, message.txid, false, op or "?", "unknown op") end
      end
    elseif ev[1] == "timer" then
      reapStale()
    end
    os.startTimer(5)
  end
end

local ok, err = pcall(main)
if not ok then
  print("Error: "..tostring(err))
  for _, tx in pairs(xfers) do cleanupTx(tx) end
end