-- Doggy OS Boot Manager for CC Tweaked

-- List of allowed disk IDs (you can populate this with the IDs of the allowed disks)
local allowedDisks = {
    "/disk2",  -- Example: disk2 is allowed by default
    "/disk3",  -- Example: disk3 is allowed by default
}

local function findStartupFiles()
    local disks = {}
    local diskPaths = {"/disk2", "/disk3"}  -- Include only /disk2/ and /disk3/

    -- Search for startup files in specified disk paths
    for _, diskPath in ipairs(diskPaths) do
        local paths = {
            diskPath .. "/startup",
            diskPath .. "/startup.lua",
        }
        for _, path in ipairs(paths) do
            if fs.exists(path) then
                table.insert(disks, {path = path, displayName = path, id = diskPath})
                break
            end
        end
    end

    -- Check /disk/ for /boot/error only (default entry, no secure boot check)
    if fs.exists("/disk/boot/error") then
        table.insert(disks, {path = "/disk/boot/error", displayName = "Doggy OS", id = "/disk"})
    end

    -- Add built-in recovery option (default entry, no secure boot check)
    if fs.exists("/disk/boot/Recovery.lua") then
        table.insert(disks, {path = "/disk/boot/Recovery.lua", displayName = "Doggy OS Recovery Environment", id = "/disk"})
    end

    return disks
end

local function isDiskAllowed(diskID)
    -- For external disks, check if the disk is allowed by secure boot
    local secureBootFile = diskID .. "/.secureboot/" .. diskID:sub(6) .. ".allowed"
    if fs.exists(secureBootFile) then
        return true
    end
    return false
end

local function drawMenu(disks, selected, bootingMessage)
    local width, height = term.getSize()
    local menuWidth, menuHeight = 40, 10 + #disks
    local startX = math.floor((width - menuWidth) / 2)
    local startY = math.floor((height - menuHeight) / 2)

    term.setBackgroundColor(colors.black)
    term.clear()

    -- Draw centered window background
    term.setBackgroundColor(colors.gray)
    for y = 0, menuHeight do
        term.setCursorPos(startX, startY + y)
        term.write(string.rep(" ", menuWidth))
    end

    -- Draw Header
    term.setCursorPos(startX + 1, startY + 1)
    term.setTextColor(colors.blue)
    term.write("  ***** Doggy OS Boot Manager *****  ")

    -- Draw Instructions
    term.setCursorPos(startX + 1, startY + 3)
    term.setTextColor(colors.lightGray)
    term.write("Select a boot device")

    term.setCursorPos(startX + 1, startY + 4)
    term.write(string.rep("=", menuWidth - 2))

    -- Draw boot options
    for i, disk in ipairs(disks) do
        term.setCursorPos(startX + 2, startY + 5 + i)
        if i == selected then
            term.setTextColor(colors.white)
            term.setBackgroundColor(colors.blue)
        else
            term.setTextColor(colors.lightGray)
            term.setBackgroundColor(colors.gray)
        end
        term.write(" [" .. i .. "] " .. disk.displayName .. string.rep(" ", menuWidth - 8 - #disk.displayName))
    end

    -- Draw Footer
    term.setCursorPos(startX + 1, startY + 6 + #disks)
    term.setTextColor(colors.lightGray)
    term.setBackgroundColor(colors.gray)
    term.write(string.rep("=", menuWidth - 2))

    -- Display booting message if available
    if bootingMessage then
        term.setCursorPos(startX + 2, startY + 7 + #disks)
        term.setTextColor(colors.cyan)
        term.write(bootingMessage)
    end

    -- Reset colors
    term.setTextColor(colors.white)
    term.setBackgroundColor(colors.black)
end

local function main()
    local disks
    while true do
        disks = findStartupFiles()
        if #disks > 0 then
            break
        end
        -- No boot files found, keep searching
        term.clear()
        term.setCursorPos(1, 1)
        term.setTextColor(colors.red)
        term.setBackgroundColor(colors.black)
        term.write("No boot files found. Searching...")
        sleep(1)
    end

    local selected = 1
    drawMenu(disks, selected, nil)

    while true do
        local event, key = os.pullEvent("key")
        if key == keys.up then
            selected = (selected - 2) % #disks + 1
        elseif key == keys.down then
            selected = selected % #disks + 1
        elseif key == keys.enter then
            -- Check if the disk is allowed to boot
            local selectedDisk = disks[selected]
            if selectedDisk.id == "/disk" or isDiskAllowed(selectedDisk.id) then
                -- Update the GUI to show booting message
                drawMenu(disks, selected, "Booting from " .. selectedDisk.path .. "...")
                sleep(1)  -- Wait for a second to display the booting message

                -- Set text color to white and background color to black before booting
                term.setTextColor(colors.white)
                term.setBackgroundColor(colors.black)

                -- Boot the selected OS
                shell.run(selectedDisk.path)
                break
            else
                -- Show error message if the disk is not allowed
                drawMenu(disks, selected, "Error: Disk ID not allowed!")
                sleep(2)  -- Display the error for 2 seconds
                drawMenu(disks, selected, nil)  -- Refresh the menu
            end
        end
        drawMenu(disks, selected, nil)  -- Refresh the menu without booting message
    end
end

main()