-- Doggy OS Systems
-- DogRE 13.0 (N3K0 V13 CFG)

local function clearScreen(message)
    term.clear()
    term.setCursorPos(1, 1)
    if message then
        print(message)
    end
end

local function waitForDisk(diskPath, message)
    clearScreen(message)
    while not fs.exists(diskPath) do
        sleep(1)
    end
end

local function promptForEnter(message)
    print(message .. " Press Enter to continue...")
    io.read()  -- Wait for user to press Enter
end

local function askToOverride(filePath)
    clearScreen("File " .. filePath .. " already exists.\nDo you want to override it?\n1. Yes\n2. No (Skip)")
    local choice = tonumber(io.read())
    if choice == 1 then
        return true  -- User wants to override
    else
        return false  -- User wants to skip
    end
end

-- Enhanced error display
local function displayError(errorMsg, errorCode)
    clearScreen("===== DOGGY OS SYSTEM RECOVERY ERROR =====")
    print("Error Message: " .. errorMsg)
    print("Error Code: DOSx7002")

    print("\nPossible Causes:")
    print("1. Missing or corrupted files.")
    print("2. Disk access errors.")
    print("3. Unexpected system behavior.")
    print("\nSuggested Actions:")
    print("1. Check disk connections and try again.")
    print("2. Run a full system recovery.")
    print("3. Contact support if the issue persists.\n")
    promptForEnter("Press Enter to Reboot system and return to normal system environment")
end

local function performAction(action, successMessage, failureMessage)
    clearScreen()
    local success, err = pcall(action)
    if success then
        print(successMessage)
    else
        -- Generate an error screen with detailed information
        displayError(failureMessage or "An unexpected error occurred.", err)
    end
    sleep(2)
end

-- Function to create a snapshot of the current system
local function createSnapshot()
    clearScreen("===== Doggy OS Snapshot =====")
    print("Enter a name for this snapshot:")
    local snapshotName = io.read()
    local snapshotDir = "./bkps/" .. snapshotName

    if fs.exists(snapshotDir) then
        if not askToOverride(snapshotDir) then
            return
        end
    end

    clearScreen("Saving snapshot...")

    -- Create the snapshot directory
    fs.makeDir(snapshotDir)

    -- Backup system files excluding /disk/users/ and /disk/config/
    local function backupDirectory(sourceDir, targetDir)
        for _, file in pairs(fs.list(sourceDir)) do
            local filePath = sourceDir .. "/" .. file
            local targetPath = targetDir .. "/" .. file
            if fs.isDir(filePath) then
                fs.makeDir(targetPath)
                backupDirectory(filePath, targetPath)
            else
                local fileData = fs.open(filePath, "r").readAll()
                local newFile = fs.open(targetPath, "w")
                newFile.write(fileData)
                newFile.close()
            end
        end
    end

    backupDirectory("/disk", snapshotDir)  -- Backup the whole /disk except the user data

    print("Snapshot saved successfully!")
    promptForEnter("Press Enter to return to the menu")
end

-- Function to list all saved snapshots
local function listSnapshots()
    local snapshots = fs.list("./bkps")
    local snapshotNames = {}

    for _, snapshot in pairs(snapshots) do
        local snapshotPath = "./bkps/" .. snapshot
        if fs.isDir(snapshotPath) then
            table.insert(snapshotNames, snapshot)
        end
    end

    if #snapshotNames == 0 then
        clearScreen("No snapshots found.")
        promptForEnter("Press Enter to return to the menu")
        return nil
    end

    clearScreen("Available Snapshots:\n")
    for i, snapshotName in ipairs(snapshotNames) do
        print(i .. ". " .. snapshotName)
    end
    print("\nEnter the number of the snapshot you want to load:")
    local choice = tonumber(io.read())

    if choice and choice >= 1 and choice <= #snapshotNames then
        return snapshotNames[choice]
    else
        clearScreen("Invalid choice. Please try again.")
        sleep(2)
        return nil
    end
end

-- Function to load a snapshot
local function loadSnapshot()
    clearScreen("===== Load Snapshot =====")
    local snapshotName = listSnapshots()

    if not snapshotName then
        return
    end

    local snapshotDir = "./bkps/" .. snapshotName
    clearScreen("Loading snapshot...")

    -- Backup user data (excluding /disk/users/ and /disk/config/)
    local function backupUserData()
        local userBackupDir = "./bkps/temp_users"
        fs.makeDir(userBackupDir)

        -- Backup /disk/users/
        local function backupDirectory(sourceDir, targetDir)
            for _, file in pairs(fs.list(sourceDir)) do
                local filePath = sourceDir .. "/" .. file
                local targetPath = targetDir .. "/" .. file
                if fs.isDir(filePath) then
                    fs.makeDir(targetPath)
                    backupDirectory(filePath, targetPath)
                else
                    local fileData = fs.open(filePath, "r").readAll()
                    local newFile = fs.open(targetPath, "w")
                    newFile.write(fileData)
                    newFile.close()
                end
            end
        end

        backupDirectory("/disk/users", userBackupDir)
        backupDirectory("/disk/config", "./bkps/temp_config")
    end

    backupUserData()

    -- Restore snapshot system files
    local function restoreDirectory(sourceDir, targetDir)
        for _, file in pairs(fs.list(sourceDir)) do
            local filePath = sourceDir .. "/" .. file
            local targetPath = targetDir .. "/" .. file
            if fs.isDir(filePath) then
                if file ~= "users" and file ~= "config" then  -- Skip user data and config data
                    fs.makeDir(targetPath)
                    restoreDirectory(filePath, targetPath)
                end
            else
                local fileData = fs.open(filePath, "r").readAll()
                local newFile = fs.open(targetPath, "w")
                newFile.write(fileData)
                newFile.close()
            end
        end
    end

    restoreDirectory(snapshotDir, "/disk")  -- Restore system files

    -- Restore user data
    local function restoreUserData()
        local function restoreDirectory(sourceDir, targetDir)
            for _, file in pairs(fs.list(sourceDir)) do
                local filePath = sourceDir .. "/" .. file
                local targetPath = targetDir .. "/" .. file
                if fs.isDir(filePath) then
                    fs.makeDir(targetPath)
                    restoreDirectory(filePath, targetPath)
                else
                    local fileData = fs.open(filePath, "r").readAll()
                    local newFile = fs.open(targetPath, "w")
                    newFile.write(fileData)
                    newFile.close()
                end
            end
        end

        restoreDirectory("./bkps/temp_users", "/disk/users")
        restoreDirectory("./bkps/temp_config", "/disk/config")
    end

    restoreUserData()

    print("Snapshot loaded and user data restored successfully!")
    promptForEnter("Press Enter to return to the menu")
end

local function recoverPassword()
    clearScreen("===== Doggy OS Password Recovery =====")
    print("Enter your username:")
    local username = io.read()
    local userPath = "/disk/users/" .. username .. "/password.txt"

    -- Prompt for the system-wide recovery key
    local recoveryKeyPath = "/recovery/rckey.txt"
    if fs.exists(recoveryKeyPath) then
        print("Enter recovery key to proceed:")
        local enteredKey = io.read()

        -- Check if the entered key matches the saved key
        local recoveryKeyFile = fs.open(recoveryKeyPath, "r")
        local storedKey = recoveryKeyFile.readAll()
        recoveryKeyFile.close()

        if enteredKey == storedKey then
            if fs.exists(userPath) then
                print("A password reset has been initiated.")
                print("Enter new password:")
                local newPassword = io.read()

                local file = fs.open(userPath, "w")
                file.write(newPassword)
                file.close()

                print("Password reset successful.")
            else
                print("User not found.")
            end
        else
            print("Invalid recovery key. Access denied.")
        end
    else
        print("Recovery key file not found.")
    end

    promptForEnter("Press Enter to return to the menu")
end

local function recoveryMenu()
    while true do
        clearScreen("===== Recovery Options =====\n1. Full System Recovery\n2. Wipe data / Reset to factory defaults\n3. Password Recovery\n4. Create Snapshot\n5. Load Snapshot\n6. Exit")
        local choice = tonumber(io.read())
        if choice == 1 then
            fullSystemRecovery()
        elseif choice == 2 then
            wipeDataAndResetFactoryDefaults()
        elseif choice == 3 then
            recoverPassword()
        elseif choice == 4 then
            createSnapshot()
        elseif choice == 5 then
            loadSnapshot()
        elseif choice == 6 then
            break
        else
            clearScreen("Invalid choice. Please try again.")
            sleep(2)
        end
    end
end

local function powerMenu()
    clearScreen("===== Power Options =====\n1. Reboot system\n2. Shutdown system")
    local choice = tonumber(io.read())
    if choice == 1 then
        os.reboot()
    elseif choice == 2 then
        os.shutdown()
    else
        clearScreen("Invalid choice. Please try again.")
        sleep(2)
    end
end

local function mainMenu()
    while true do
        clearScreen("===== Doggy OS DEV Recovery GUI =====\n1. Recovery Options\n2. Power Options\n3. Exit")
        local choice = tonumber(io.read())
        if choice == 1 then
            recoveryMenu()
        elseif choice == 2 then
            powerMenu()
        else
            clearScreen("Invalid choice. Please try again.")
            sleep(2)
        end
    end
end

-- Start the main menu
if fs.exists("/disk/config/dogRE/disableRE.cfg/") then
    clearScreen("===== DOGGY OS SYSTEM RECOVERY ERROR =====")
    print("Error Message: " .. "Recovery Environment Disabled")
    print("Error Code: DOSx7000")
    print("\nPossible Causes:")
    print("1. Recovery Environment has been disabled by your system administrator")
    print("2. Malware or unwanted software disabled DogRE")
    print("3. System config errors or corruption")
    print("\nSuggested Actions:")
    print("1. Reboot the computer and re-enable DogRE")
    print("2. Reinstall Doggy OS or rebuild the recovery partition")
    print("3. Contact support if the issue persists.\n")
    promptForEnter("Press Enter to Reboot system")
else
    mainMenu()
end