#!/usr/bin/lua

-- Required modules
require "ubus"
require "luci.util"
local json = require "luci.jsonc"
local nixio = require "nixio"

-- Constants
local LOGIN_PASSWORD = "admin123"
local WIFI_IFACE_NAMES = {
    radio0 = "apcli24g",
    radio1 = "apcli5g"
}

-- Utility functions
local function is_logged_in(post_data)
    return post_data and post_data["session"] == "logged_in"
end

local function match_login_password(login_password)
    -- get the password from the /etc/config/wifi_config file
    -- the wifi_config file should have the following content:
    -- config wifi_config 
    --     option password 'admin123'
    -- check if /etc/config/wifi_config file exists
    if not nixio.fs.access("/etc/config/wifi_config") then
        return true
    end
    local password = (luci.util.exec("uci -q get wifi_config.@wifi_config[0].password") or ""):gsub("^%s*(.-)%s*$", "%1")
    return login_password == password
end

local function parse_post_data()
    local post_data = nil
    local content_length = tonumber(os.getenv("CONTENT_LENGTH")) or 0
    if content_length > 0 then
        local raw_data = nixio.stdin:read(content_length)
        post_data = {}
        for key, value in string.gmatch(raw_data, "([^&=]+)=([^&=]*)&?") do
            post_data[key] = value
        end
    end
    return post_data
end

-- WiFi configuration functions
local function execute_uci_commands(commands)
    for _, cmd in ipairs(commands) do
        os.execute("uci -q " .. cmd)
    end
    os.execute("uci -q commit wireless")
    os.execute("uci -q commit network")
    os.execute("uci -q commit firewall")
    -- Before reboot, notify the user that the device will reboot
    print("<h1>Device will reboot</h>")
    print("<p>Please wait for your router restarting</p>")
    os.execute("reboot")
end

local function configure_wifi(ssid, password, radio, encryption)
    if not (ssid and password and radio and encryption) then return end

    -- URL decode ssid and password
    ssid = luci.util.urldecode(ssid)
    password = luci.util.urldecode(password)

    local wifi_iface_name = WIFI_IFACE_NAMES[radio]
    local commands = {
        "delete network.wwan",
        "add network interface", 
        "set network.@interface[-1].proto='dhcp'",
        "rename network.@interface[-1]='wwan'",
        "del_list firewall.@zone[1].network='wwan'",
        "add_list firewall.@zone[1].network='wwan'",
        "delete wireless.apcli24g",
        "delete wireless.apcli5g",
        "add wireless wifi-iface",
        string.format("set wireless.@wifi-iface[-1].device='%s'", radio),
        "set wireless.@wifi-iface[-1].mode='sta'",
        "set wireless.@wifi-iface[-1].network='wwan'",
        string.format("set wireless.@wifi-iface[-1].ssid='%s'", ssid),
        string.format("set wireless.@wifi-iface[-1].encryption='%s'", encryption),
        string.format("set wireless.@wifi-iface[-1].key='%s'", password),
        "set wireless.@wifi-iface[-1].disabled='0'",
        string.format("rename wireless.@wifi-iface[-1]='%s'", wifi_iface_name),
    }
    execute_uci_commands(commands)
end

-- Network scanning functions
local function get_network_info(conn, radio)
    local results = conn:call("iwinfo", "scan", { device = radio })
    local networks = {}
    if results and results.results then
        for _, network in ipairs(results.results) do
            if network.ssid then
                local auth = "none"
                local enc = (type(network.encryption) == "table") and network.encryption or nil
                local is_wep = enc and type(enc.wep) == "table"
                local is_psk = enc and type(enc.wpa) == "table" and luci.util.contains(enc.authentication, 'psk')
                local is_sae = enc and type(enc.wpa) == "table" and luci.util.contains(enc.authentication, 'sae')
                
                if is_sae then
                    auth = "sae"
                elseif is_psk then
                    for i = #enc.wpa, 1, -1 do
                        if enc.wpa[i] == 2 then
                            auth = "psk2"
                            break
                        elseif enc.wpa[i] == 1 then
                            auth = "psk"
                            break
                        end
                    end
                elseif is_wep then
                    auth = "wep"
                else 
                    auth = "none"
                end
                
                table.insert(networks, {
                    ssid = network.ssid,
                    bssid = network.bssid,
                    signal = network.signal,
                    encryption = auth,
                    encryption_info = network.authentication
                })
            end
        end
        -- Sort networks by signal strength (higher values first)
        table.sort(networks, function(a, b) return a.signal > b.signal end)
    end
    return networks
end

-- HTML rendering functions
local function render_network_options(networks)
    local html = '<option value="none" data-auth="none">none</option>'
    for _, network in ipairs(networks) do
        html = html .. string.format(
            '<option value="%s" data-auth="%s">%s (BSSID: %s, Signal: %ddBm)</option>',
            network.ssid,
            network.encryption,
            network.ssid,
            network.bssid,
            network.signal
        )
    end
    return html
end

local function show_login_form(error_message)
    print([[<h1>WiFi Configuration</h1>]])
    print([[<div class="form-container">]])
    if error_message then
        print([[<p style="color: red;">]] .. error_message .. [[</p>]])
    end
    print([[<form method="post" action="/cgi-bin/wifi-config">]])
    print([[<input type="password" name="login_password" placeholder="Password">]])
    print([[<input type="submit" value="Login">]])
    print([[</form>]])
    print([[</div>]])
end

-- HTML template functions
local function render_wifi_toggle()
    return [[
        <div class="wifi-toggle">
            <input type="checkbox" id="wifiToggle" onchange="toggleWifi()">
            <label for="wifiToggle">Show 2.4GHz Networks</label>
        </div>
    ]]
end

local function render_network_section(band, networks, is_hidden)
    local html = string.format([[
        <div id="networks%s" style="display: %s;">
            <h2>%s Networks</h2>
            <select id="networks%sSelect" onchange="showEncryption('%s', this.selectedIndex)">
            %s
            </select>
            <input type="text" id="auth%s" readonly>
        </div>
    ]], 
    band,
    is_hidden and "none" or "block",
    band == "2G" and "2.4GHz" or "5GHz",
    band,
    band,
    render_network_options(networks),
    band)
    return html
end

local function render_javascript()
    return [[
        <script>
        function toggleWifi() {
            var networks2G = document.getElementById('networks2G');
            var networks5G = document.getElementById('networks5G');
            var checkbox = document.getElementById('wifiToggle');
            networks2G.style.display = checkbox.checked ? 'block' : 'none';
            networks5G.style.display = checkbox.checked ? 'none' : 'block';
        }

        function showEncryption(band, index) {
            var select = document.getElementById('networks' + band + 'Select');
            var auth = document.getElementById('auth' + band);
            if (select.options[index]) {
                auth.value = select.options[index].getAttribute('data-auth') || 'none';
            }
        }

        function prepareSubmission() {
            var is2G = document.getElementById('wifiToggle').checked;
            var select = document.getElementById(is2G ? 'networks2GSelect' : 'networks5GSelect');
            var auth = document.getElementById(is2G ? 'auth2G' : 'auth5G');
            var password = document.getElementById('wifiPassword');

            if (!select || !select.value || select.value === 'none' || select.value === '') {
                alert('Please select a valid network');
                return false;
            }

            // Only require password if authentication is not 'none'
            if (auth.value !== 'none' && !password.value) {
                alert('Password is required for encrypted networks');
                return false;
            } 

            if (!confirm('This operation will cause the device to reboot. Do you want to continue?')) {
                return false;
            }

            document.getElementById('selectedSSID').value = select.value;
            document.getElementById('selectedEncryption').value = auth.value;
            document.getElementById('selectedRadio').value = is2G ? 'radio0' : 'radio1';

            return true;
        }
        </script>
    ]]
end

local function show_wifi_form(conn)
    print([[
        <div class="form-container" id="wifiFormContainer">
            <h1>WiFi Configuration</h1>
            <form id="wifiForm" action="/cgi-bin/wifi-config" method="post">
                <input type="hidden" name="ssid" id="selectedSSID">
                <input type="hidden" name="encryption" id="selectedEncryption">
                <input type="hidden" name="radio" id="selectedRadio">
                <input type="hidden" name="session" value="logged_in">
    ]])

    print(render_wifi_toggle())
    print(render_network_section("2G", get_network_info(conn, "radio0"), true))
    print(render_network_section("5G", get_network_info(conn, "radio1"), false))

    print([[
                <input type="password" name="password" id="wifiPassword" placeholder="Enter WiFi Password">
                <input type="submit" value="Connect" onclick="return prepareSubmission()">
            </form>
        </div>
    ]])

    print(render_javascript())
end

-- Main execution
local function main()
    -- Connect to ubus
    local conn = ubus.connect()
    if not conn then
        error("Failed to connect to ubus")
    end

    -- Write HTTP headers and CSS
    print("Content-Type: text/html\n")
    print([[<!DOCTYPE html><html><head><title>WiFi Configuration</title>]])
    print([[
        <style>
            body {
            font-family: Arial, sans-serif;
            max-width: 800px;
            margin: 20px auto;
            padding: 20px;
            background-color: #f5f5f5;
            }
            h1 {
                color: #333;
                text-align: center;
                margin-bottom: 30px;
            }
            h2 {
                color: #444;
                margin-top: 20px;
            }
            .form-container {
                background: white;
                padding: 20px;
                border-radius: 8px;
                box-shadow: 0 2px 4px rgba(0,0,0,0.1);
            }
            .wifi-toggle {
                margin: 20px 0;
            }
            .wifi-toggle label {
                font-size: 16px;
                margin-left: 8px;
                cursor: pointer;
            }
            select {
                width: 100%;
                padding: 10px;
                margin: 10px 0;
                border: 1px solid #ddd;
                border-radius: 4px;
                font-size: 14px;
            }
            input[type="text"] {
                width: 100%;
                padding: 10px;
                margin: 10px 0;
                border: 1px solid #ddd;
                border-radius: 4px;
                font-size: 14px;
            }
            input[type="password"] {
                width: 100%;
                padding: 10px;
                margin: 10px 0;
                border: 1px solid #ddd;
                border-radius: 4px;
                font-size: 14px;
            }
            input[type="submit"] {
                width: 100%;
                padding: 12px;
                background-color: #4CAF50;
                color: white;
                border: none;
                border-radius: 4px;
                cursor: pointer;
                font-size: 16px;
                margin-top: 20px;
            }
            input[type="submit"]:hover {
                background-color: #45a049;
            }
        </style>
    ]])
    print([[</head><body>]])

    -- Handle requests
    local post_data = parse_post_data()
    if post_data then
        if post_data["login_password"] then
            local login_password = post_data["login_password"]
            if match_login_password(login_password) then
                show_wifi_form(conn)
            else
                show_login_form("Invalid password")
            end
        elseif is_logged_in(post_data) then
            configure_wifi(
                post_data["ssid"],
                post_data["password"],
                post_data["radio"],
                post_data["encryption"]
            )
        else
            show_login_form()
        end
    else
        show_login_form()
    end

    print([[</body></html>]])
    conn:close()
end

-- Start the application
main()