-- Function to decrypt text using a key
function decrypt(text, key)
    local map = invertMap(createMap(key))
    local decrypted = {}
    for i = 1, #text do
        local char = text:sub(i, i)
        table.insert(decrypted, map[char] or char)
    end
    return table.concat(decrypted)
end

-- Function to load the key from a file and check if decryption is possible
function checkAndDecrypt()
    local rootFiles = fs.list("/")
    local keyFile = nil

    -- Check for key file in the root directory
    for _, file in ipairs(rootFiles) do
        if string.match(file, "^key%d+%.txt$") then
            keyFile = file
            break
        end
    end

    if keyFile then
        showProtectionScreen(keyFile)
    else
        print("No encryption key found. Use 'encrypt' mode to encrypt files.")
    end
end

-- Function to show protection screen options
function showProtectionScreen(keyFile)
    term.clear()
    term.setCursorPos(1, 1)
    print("Doggy OS File System Protection")
    print()
    print("F1 - Unlock")
    print("F8 - Shutdown")
    print("F9 - Recover System")

    while true do
        local event, key = os.pullEvent("key")
        if key == keys.f1 then
            local key = loadKeyFromFile(keyFile)
            processDirectory("/disk", decrypt, key)
            shell.run("/disk/boot/start-check")
            return
        elseif key == keys.f8 then
            os.shutdown()
            return
        elseif key == keys.f9 then
            recoverSystem()
            return
        end
    end
end

-- Function to recover the system (delete disk and replace with recovery)
function recoverSystem()
    local files = fs.list("/disk")
    for _, file in ipairs(files) do
        local filePath = fs.combine("/disk", file)
        if not fs.isDir(filePath) then
            fs.delete(filePath)
        end
    end
    -- Optionally copy recovery files from a backup location or setup
    print("System recovered. Please restart your system.")
end

-- Main Program
term.clear()
term.setCursorPos(1, 1)
print("Doggy OS File System Encryption")
print()
print("Initializing...")

local tArgs = {...}
local mode = tArgs[1]

if mode == "encrypt" then
    local key = generateKey()
    encryptDirectory("/disk", key)
    local keyFileName = saveKeyToFile(key)
    print("All files in /disk/ encrypted.")
    print("Key saved to file: " .. keyFileName)
    print()
    print("Encryption complete.")
else
    print("Doggy OS File System Protection")
    print()
    checkAndDecrypt()
end