ndisk_api.lua

-- ndisk_api.lua
-- API for Distributed NDisk System
-- RAID10 Edition: Each chunk stored on primary + mirror node
-- UPDATED: Robustness, Security, and Recovery Features

local ndisk = {}

local modem = peripheral.find("modem")
if not modem then error("No modem found") end
modem.open(5) -- Channel 5 for NDisk storage traffic

local nodes = {}
local chunkMap = {}
local initialized = false
local systemSecret = nil -- Use ndisk.setSecret()
local writeLocks = {}    -- Active write locks per path

-- ============================================
-- RAID10 HELPER FUNCTIONS
-- ============================================

-- Get mirror node index (offset by half the node count for distribution)
local function getMirrorIndex(primaryIdx, nodeCount)
    if nodeCount < 2 then return nil end
    local offset = math.floor(nodeCount / 2)
    return ((primaryIdx - 1 + offset) % nodeCount) + 1
end

-- Get sorted node list for consistent ordering
local function getSortedNodes()
    local sorted = {}
    for idx, id in pairs(nodes) do
        table.insert(sorted, {idx = idx, id = id})
    end
    table.sort(sorted, function(a, b) return a.idx < b.idx end)
    return sorted
end

-- Path Sanitization (Security)
local function checkPath(path)
    if not path or type(path) ~= "string" then return false, "Invalid path" end
    if path:find("%.%.") then return false, "Invalid path: '..' not allowed" end
    if path:find(":") then return false, "Invalid path: ':' not allowed" end
    return true
end

-- ============================================
-- NODE PERSISTENCE
-- ============================================

local function saveNodeIds()
    local file = fs.open("ids.txt", "w")
    for hierarchicalNum, nodeId in pairs(nodes) do
        file.writeLine(nodeId .. " " .. hierarchicalNum)
    end
    file.close()
end

local function loadNodeIds()
    nodes = {}
    if fs.exists("ids.txt") then
        local file = fs.open("ids.txt", "r")
        for line in file.readLine do
            local nodeId, hierarchicalNum = line:match("(%d+)%s+(%d+)")
            if nodeId and hierarchicalNum then
                nodes[tonumber(hierarchicalNum)] = tonumber(nodeId)
            end
        end
        file.close()
    end
end

-- ============================================
-- CHUNK MAP PERSISTENCE (RAID10 FORMAT)
-- ============================================

local function saveChunkMap()
    local file = fs.open("chunklocation.txt", "w")
    for filename, chunks in pairs(chunkMap) do
        for chunkNum, info in pairs(chunks) do
            if type(info) == "table" then
                -- RAID10 format: path chunk# primary mirror [degraded]
                local mirrorId = info.mirror or info.primary
                local line = filename .. " chunk" .. chunkNum .. " " .. info.primary .. " " .. mirrorId
                if info.degraded then
                    line = line .. " degraded"
                end
                file.writeLine(line)
            else
                -- Legacy single-node format (backward compat)
                file.writeLine(filename .. " chunk" .. chunkNum .. " " .. info)
            end
        end
    end
    file.close()
end

local function loadChunkMap()
    chunkMap = {}
    if fs.exists("chunklocation.txt") then
        local file = fs.open("chunklocation.txt", "r")
        for line in file.readLine do
            -- Parse line into parts
            local parts = {}
            for part in line:gmatch("%S+") do
                table.insert(parts, part)
            end

            if #parts >= 3 then
                local filename = parts[1]
                local chunkMatch = parts[2]:match("chunk(%d+)")
                local chunkNum = chunkMatch and tonumber(chunkMatch) or nil
                local primaryId = tonumber(parts[3])
                local mirrorId = parts[4] and tonumber(parts[4]) or nil

                -- Skip malformed lines
                if filename and chunkNum and primaryId then
                    chunkMap[filename] = chunkMap[filename] or {}

                    -- Check for degraded flag (5th part)
                    local isDegraded = parts[5] == "degraded"

                    -- Store as RAID10 format if mirror exists and differs
                    if mirrorId and mirrorId ~= primaryId then
                        chunkMap[filename][chunkNum] = {primary = primaryId, mirror = mirrorId, degraded = isDegraded}
                    else
                        -- Legacy format or single-node: store as table for consistency
                        chunkMap[filename][chunkNum] = {primary = primaryId, mirror = nil, degraded = isDegraded}
                    end
                end
            end
        end
        file.close()
    end
end

-- ============================================
-- WRITE JOURNALING (Recovery)
-- ============================================

local function logJournal(action, path)
    local f = fs.open("ndisk_journal.txt", "a")
    f.writeLine(os.epoch("utc") .. "|" .. action .. "|" .. path)
    f.close()
end

local function clearJournal()
    if fs.exists("ndisk_journal.txt") then
        fs.delete("ndisk_journal.txt")
    end
end

local function checkJournal()
    if fs.exists("ndisk_journal.txt") then
        term.setTextColor(colors.orange)
        print("NDisk Alert: Previous session terminated abnormally.")
        print("Checking journal for incomplete writes...")
        term.setTextColor(colors.white)

        -- Simple recovery: just warn for now, real rollback is complex
        local f = fs.open("ndisk_journal.txt", "r")
        local content = f.readAll()
        f.close()
        print("Found pending operations:\n" .. content)
        print("Note: File consistency may need manual verification.")

        -- Back up the journal and clear active one
        local backup = fs.open("ndisk_journal_old.txt", "w")
        backup.write(content)
        backup.close()
        fs.delete("ndisk_journal.txt")
    end
end

-- ============================================
-- NODE DISCOVERY
-- ============================================

local function discoverNodes()
    print("NDisk: Discovering storage nodes...")
    -- Discovery sends secret if available, but nodes might not require it for discovery?
    -- Actually, safer to include it if we have it
    modem.transmit(5, 5, {type = "discovery", secret = systemSecret})
    local timer = os.startTimer(2)
    local nextHierarchicalNum = 1
    for num, _ in pairs(nodes) do
        if num >= nextHierarchicalNum then
            nextHierarchicalNum = num + 1
        end
    end

    while true do
        local event, param1, param2, param3, param4 = os.pullEvent()
        if event == "modem_message" and type(param4) == "table" and param4.type == "discovery_response" then
            local nodeId = param4.id
            local hierarchicalNum = nil
            for num, id in pairs(nodes) do
                if id == nodeId then
                    hierarchicalNum = num
                    break
                end
            end
            if not hierarchicalNum then
                hierarchicalNum = nextHierarchicalNum
                nodes[hierarchicalNum] = nodeId
                nextHierarchicalNum = nextHierarchicalNum + 1
            end
        elseif event == "timer" and param1 == timer then
            break
        end
    end
    saveNodeIds()
    return nodes
end

-- ============================================
-- NODE COMMUNICATION (WITH RAID10 FALLBACK)
-- ============================================

local function generateRequestId()
    return os.getComputerID() .. "_" .. os.clock() .. "_" .. math.random(10000)
end

local function sendRequest(nodeId, request, fallbackNodeId)
    request.targetNode = nodeId

    -- AUTHENTICATION: Attach secret
    if systemSecret then
        request.secret = systemSecret
    end

    modem.transmit(5, 5, request)
    local timeout = 2.0  -- UPDATED: Increased from 0.5s for robustness
    local timer = os.startTimer(timeout)

    while true do
        local event, p1, p2, p3, p4 = os.pullEvent()
        if event == "modem_message" and type(p4) == "table" and p4.id == request.id then
            return p4.response
        elseif event == "timer" and p1 == timer then
            -- Primary timed out, try fallback if available
            if fallbackNodeId and fallbackNodeId ~= nodeId then
                request.targetNode = fallbackNodeId
                modem.transmit(5, 5, request)
                local timer2 = os.startTimer(timeout)  -- Fallback also gets more time

                while true do
                    local ev2, pp1, pp2, pp3, pp4 = os.pullEvent()
                    if ev2 == "modem_message" and type(pp4) == "table" and pp4.id == request.id then
                        return pp4.response
                    elseif ev2 == "timer" and pp1 == timer2 then
                        return {success = false, error = "Both primary and mirror timeout"}
                    end
                end
            end
            return {success = false, error = "Timeout waiting for Node " .. nodeId}
        end
    end
end

-- ============================================
-- API FUNCTIONS
-- ============================================

function ndisk.setSecret(secret)
    systemSecret = secret
end

function ndisk.init()
    if initialized then return end

    -- Journal Check
    checkJournal()

    loadNodeIds()
    chunkMap = {}
    loadChunkMap()

    if not next(nodes) then
        discoverNodes()
    end

    local nodeCount = 0
    for _ in pairs(nodes) do nodeCount = nodeCount + 1 end

    if nodeCount == 0 then
        print("NDisk Warning: No storage nodes found/configured.")
    else
        local raidStatus = nodeCount >= 2 and "RAID10 enabled" or "Single-node mode"
        print("NDisk: Initialized with " .. nodeCount .. " nodes. (" .. raidStatus .. ")")
    end
    initialized = true
end

function ndisk.refreshNodes()
    discoverNodes()
end

function ndisk.write(path, content)
    if not initialized then ndisk.init() end

    local valid, err = checkPath(path)
    if not valid then return false, err end

    -- FILE-LEVEL LOCKING
    while writeLocks[path] do
        os.sleep(0.1)
    end
    writeLocks[path] = os.epoch("utc")

    -- Protected Call to ensure lock release
    local function doWrite()
        -- Log Intent
        logJournal("WRITE", path)

        local sortedNodes = getSortedNodes()
        local nodeCount = #sortedNodes

        if nodeCount == 0 then
            discoverNodes()
            sortedNodes = getSortedNodes()
            nodeCount = #sortedNodes
            if nodeCount == 0 then return false, "No storage nodes available" end
        end

        local chunkSize = math.ceil(#content / nodeCount)
        if chunkSize == 0 then chunkSize = 1 end

        chunkMap[path] = {}

        local writeSuccess = true
        local writeError = nil

        for i, nodeInfo in ipairs(sortedNodes) do
            local start = (i - 1) * chunkSize + 1
            if start > #content and i > 1 then break end

            local chunk = content:sub(start, math.min(start + chunkSize - 1, #content))
            if #content == 0 and i == 1 then chunk = "" end

            local primaryId = nodeInfo.id
            local mirrorIdx = getMirrorIndex(i, nodeCount)
            local mirrorId = mirrorIdx and sortedNodes[mirrorIdx] and sortedNodes[mirrorIdx].id or nil

            -- Write to primary
            local primaryOk = false
            local reqId = generateRequestId()
            local response = sendRequest(primaryId, {
                id = reqId,
                action = "write",
                path = path,
                chunk = i,
                content = chunk
            })
            if response and response.success then
                primaryOk = true
            end

            -- Write to mirror (if available and different)
            local mirrorOk = false
            if mirrorId and mirrorId ~= primaryId then
                local mirrorReqId = generateRequestId()
                local mirrorResp = sendRequest(mirrorId, {
                    id = mirrorReqId,
                    action = "write",
                    path = path,
                    chunk = i,
                    content = chunk
                })
                if mirrorResp and mirrorResp.success then
                    mirrorOk = true
                end
            else
                -- No mirror means primary success is enough
                mirrorOk = primaryOk
            end

            -- RESILIENCE: Need at least ONE success, not both
            if not primaryOk and not mirrorOk then
                writeSuccess = false
                writeError = "Failed to write chunk " .. i .. " to both primary and mirror"
                break
            end

            -- Mark as degraded if only one node has the data
            local isDegraded = not (primaryOk and mirrorOk)
            chunkMap[path][i] = {primary = primaryId, mirror = mirrorId, degraded = isDegraded}

            if start + chunkSize > #content then break end
        end

        if writeSuccess then
            saveChunkMap()
            clearJournal() -- Operation complete
            return true
        else
            return false, writeError
        end
    end

    local status, res, err = pcall(doWrite)

    -- Always release lock
    writeLocks[path] = nil

    if not status then return false, "Crash during write: " .. tostring(res) end
    if not res then return false, err end
    return true
end

function ndisk.read(path)
    if not initialized then ndisk.init() end
    local valid, err = checkPath(path)
    if not valid then return nil, err end

    if not chunkMap[path] then
        return nil, "File not found"
    end

    local contentChunks = {}
    local chunkCount = 0
    for _ in pairs(chunkMap[path]) do chunkCount = chunkCount + 1 end

    for chunkNum = 1, chunkCount do
        local chunkInfo = chunkMap[path][chunkNum]
        if not chunkInfo then return nil, "Missing chunk info for chunk " .. chunkNum end

        -- Handle both RAID10 format and legacy format
        local primaryId, mirrorId
        if type(chunkInfo) == "table" then
            primaryId = chunkInfo.primary
            mirrorId = chunkInfo.mirror
        else
            primaryId = chunkInfo
            mirrorId = nil
        end

        -- Read with fallback to mirror
        local response = sendRequest(primaryId, {
            id = generateRequestId(),
            action = "read",
            path = path,
            chunk = chunkNum
        }, mirrorId)

        if response and response.success then
            contentChunks[chunkNum] = response.content or ""
        else
            return nil, "Failed to read chunk " .. chunkNum .. " from nodes"
        end
    end

    -- Ensure no nil values in contentChunks before concat
    for i = 1, chunkCount do
        if contentChunks[i] == nil then
            return nil, "Missing content for chunk " .. i
        end
    end

    return table.concat(contentChunks)
end

function ndisk.delete(path)
    if not initialized then ndisk.init() end
    local valid, err = checkPath(path)
    if not valid then return false, err end

    if not chunkMap[path] then return false, "File not found" end

    -- Log Intent
    logJournal("DELETE", path)

    for chunkNum, chunkInfo in pairs(chunkMap[path]) do
        local primaryId, mirrorId
        if type(chunkInfo) == "table" then
            primaryId = chunkInfo.primary
            mirrorId = chunkInfo.mirror
        else
            primaryId = chunkInfo
            mirrorId = nil
        end

        -- Delete from primary
        sendRequest(primaryId, {
            id = generateRequestId(),
            action = "delete",
            path = path,
            chunk = chunkNum
        })

        -- Delete from mirror
        if mirrorId and mirrorId ~= primaryId then
            sendRequest(mirrorId, {
                id = generateRequestId(),
                action = "delete",
                path = path,
                chunk = chunkNum
            })
        end
    end

    chunkMap[path] = nil
    saveChunkMap()
    clearJournal()
    return true
end

function ndisk.exists(path)
    if not initialized then ndisk.init() end
    if chunkMap[path] then return true end

    local dirPath = path .. "/"
    if path == "" then dirPath = "" end

    for filename, _ in pairs(chunkMap) do
        if filename == path then return true end
        if string.sub(filename, 1, #dirPath) == dirPath then
            return true
        end
    end
    return false
end

function ndisk.isDir(path)
    if not initialized then ndisk.init() end
    if chunkMap[path] then return false end

    local dirPath = path .. "/"
    if path == "" then return true end

    for filename, _ in pairs(chunkMap) do
        if string.sub(filename, 1, #dirPath) == dirPath then
            return true
        end
    end
    return false
end

function ndisk.list(path)
    if not initialized then ndisk.init() end
    local items = {}
    local seen = {}

    local prefix = path
    if #prefix > 0 and string.sub(prefix, -1) ~= "/" then
        prefix = prefix .. "/"
    end
    if prefix == "/" then prefix = "" end

    for filename, _ in pairs(chunkMap) do
        if string.find(filename, prefix, 1, true) == 1 then
            local remainder = string.sub(filename, #prefix + 1)
            if #remainder > 0 then
                local firstSlash = string.find(remainder, "/")
                local item
                if firstSlash then
                    item = string.sub(remainder, 1, firstSlash - 1)
                else
                    item = remainder
                end

                if not seen[item] then
                    table.insert(items, item)
                    seen[item] = true
                end
            end
        end
    end
    return items
end

function ndisk.getFreeSpace()
    if not initialized then ndisk.init() end
    local total = 0
    for _, nodeId in pairs(nodes) do
        local response = sendRequest(nodeId, {
            id = generateRequestId(),
            action = "getSpace"
        })
        if response and response.success then
            total = total + response.space
        end
    end
    -- RAID10 uses 50% for mirroring
    local nodeCount = 0
    for _ in pairs(nodes) do nodeCount = nodeCount + 1 end
    if nodeCount >= 2 then
        return math.floor(total / 2)
    end
    return total
end

-- Get RAID status info
function ndisk.getRaidStatus()
    if not initialized then ndisk.init() end
    local nodeCount = 0
    for _ in pairs(nodes) do nodeCount = nodeCount + 1 end

    return {
        enabled = nodeCount >= 2,
        nodeCount = nodeCount,
        maxFailures = math.floor(nodeCount / 2),
        mode = nodeCount >= 2 and "RAID10" or "Single"
    }
end

-- ============================================
-- BACKGROUND REPAIR
-- ============================================

function ndisk.repairDegraded()
    if not initialized then ndisk.init() end

    local degradedCount = 0
    local repairedCount = 0

    print("NDisk Repair: Scanning for degraded chunks...")

    for path, chunks in pairs(chunkMap) do
        for chunkNum, info in pairs(chunks) do
            if info.degraded then
                degradedCount = degradedCount + 1

                -- Attempt rewrite to restore redundancy
                local content = ndisk.read(path) -- This reads whole file, bit inefficient but safe
                if content then
                    print("Repairing: " .. path)
                    local success, err = ndisk.write(path, content)
                    if success then
                        repairedCount = repairedCount + 1
                    else
                        print("Failed to repair " .. path .. ": " .. (err or "unknown"))
                    end
                    -- Break inner loop to avoid re-repairing other chunks of same file needlessly if we did full write
                    break
                else
                    print("Critical: Could not read " .. path .. " to repair it")
                end
            end
        end
    end

    return {found = degradedCount, repaired = repairedCount}
end

return ndisk