-- Function to draw the title bar with "DNS Server"
local function drawTitleBar()
    term.setBackgroundColor(colors.blue)
    term.setTextColor(colors.white)
    term.setCursorPos(1, 1)
    term.write("DNS Server")
    term.setBackgroundColor(colors.black)
    term.setTextColor(colors.white)
end

-- Function to load DNS mappings from the /DNS/URLS/ directory
local function loadDNS()
    local dns = {}
    if fs.exists("/DNS/URLS") then
        for _, file in ipairs(fs.list("/DNS/URLS")) do
            local f = fs.open("/DNS/URLS/" .. file, "r")
            local ip = f.readAll()
            f.close()
            dns[file] = ip
        end
    else
        fs.makeDir("/DNS/URLS")  -- Create the directory if it doesn't exist
    end
    return dns
end

-- Function to log DNS events
local function logConnection(message)
    local logFile = fs.open("dns_log.txt", "a")
    logFile.writeLine(message)
    logFile.close()
    print(message)
end

-- Function to prompt for and verify the password
local function promptForPassword()
    term.clear()
    drawTitleBar()
    print("Enter DNS Password:")
    return read("*")  -- Mask the input for security
end

-- Function to authenticate password
local function authenticate(password)
    local filePath = "/DNS/locked.cfg"
    if fs.exists(filePath) then
        local file = fs.open(filePath, "r")
        local storedPassword = file.readAll()
        file.close()
        return password == storedPassword
    end
    return true  -- No password required if file does not exist
end

-- Function to display and manage DNS settings
local function dnsSettingsMenu()
    while true do
        term.clear()
        drawTitleBar()
        print("DNS Settings Menu:")
        print("1. View DNS Mappings")
        print("2. Manage DNS Mappings")
        print("3. Change Password")
        print("4. Back to Main Menu")

        local choice = tonumber(read())
        if choice == 1 then
            -- View DNS Mappings
            term.clear()
            drawTitleBar()
            print("Current DNS Mappings:")
            local dns = loadDNS()
            if next(dns) == nil then
                print("No DNS mappings found.")
            else
                for url, ip in pairs(dns) do
                    print(url .. " -> " .. ip)
                end
            end
            print("Press any key to return to the menu.")
            os.pullEvent("key")
        elseif choice == 2 then
            -- Manage DNS Mappings
            term.clear()
            drawTitleBar()
            print("Manage DNS Mappings:")
            print("1. Add DNS Mapping")
            print("2. Delete DNS Mapping")
            print("3. Back to DNS Settings Menu")

            local subChoice = tonumber(read())
            if subChoice == 1 then
                -- Add DNS Mapping
                term.clear()
                drawTitleBar()
                print("Enter URL to map:")
                local url = read()
                print("Enter IP address for " .. url .. ":")
                local ip = read()

                if url ~= "" and ip ~= "" then
                    local file = fs.open("/DNS/URLS/" .. url, "w")
                    file.write(ip)
                    file.close()
                    print("Mapping added successfully.")
                else
                    print("Invalid input. Mapping not added.")
                end
                print("Press any key to return to the menu.")
                os.pullEvent("key")
            elseif subChoice == 2 then
                -- Delete DNS Mapping
                term.clear()
                drawTitleBar()
                print("Enter URL to delete:")
                local url = read()

                local filePath = "/DNS/URLS/" .. url
                if fs.exists(filePath) then
                    fs.delete(filePath)
                    print("Mapping deleted successfully.")
                else
                    print("URL not found. Mapping not deleted.")
                end
                print("Press any key to return to the menu.")
                os.pullEvent("key")
            elseif subChoice == 3 then
                -- Return to DNS Settings Menu
                break
            else
                print("Invalid option. Press any key to return to the menu.")
                os.pullEvent("key")
            end
        elseif choice == 3 then
            -- Change Password
            term.clear()
            drawTitleBar()
            print("Enter new DNS password (leave blank to remove password):")
            local newPassword = read("*")
            local filePath = "/DNS/locked.cfg"
            if newPassword == "" then
                if fs.exists(filePath) then
                    fs.delete(filePath)
                    print("Password removed.")
                else
                    print("No password set.")
                end
            else
                local file = fs.open(filePath, "w")
                file.write(newPassword)
                file.close()
                print("Password updated.")
            end
            print("Press any key to return to the menu.")
            os.pullEvent("key")
        elseif choice == 4 then
            -- Return to main DNS request handling
            break
        else
            print("Invalid option. Press any key to return to the menu.")
            os.pullEvent("key")
        end
    end
end

-- Function to handle incoming DNS requests
local function handleRequests()
    rednet.open("top")  -- Assuming the modem is on the top
    print("DNS Server running... Press F1 for settings.")

    while true do
        local event, param1, param2 = os.pullEvent()

        if event == "rednet_message" then
            local senderID, url = param1, param2
            local dns = loadDNS()
            local ip = dns[url]
            logConnection("Received request for URL: " .. url)

            if not ip then
                rednet.send(senderID, "404_NOT_FOUND")
                logConnection("URL not found: " .. url)
            else
                rednet.send(senderID, ip)
                logConnection("URL resolved to IP: " .. ip)
            end

        elseif event == "key" and param1 == keys.f1 then
            -- Open DNS settings menu when F1 is pressed
            dnsSettingsMenu()
            -- Redisplay server status after returning from the menu
            term.clear()
            drawTitleBar()
            print("DNS Server running... Press F1 for settings.")
        end
    end
end

-- Main function to run the DNS server
local function main()
    handleRequests()
end

main()