-- lemonos3_fixed.lua
-- LemonOS 3.0 (fixed) - integrated chat + DM + multi-server relay client
-- Reworked to avoid rednet.lookup single-number return bug and other stability fixes.

-- CONFIG / PATHS
local CONFIG_PATH = "/lemonos.cfg"
local STARTUP_PATH = "/startup.lua"
local firmware = { device = "Lemonphone", version = "LemonOS 3.0 (fixed)" }

-- SETTINGS (persisted)
local config = {
    username = "User",
    accentColor = colors.yellow,
    bgColor = colors.black,
    rednetSide = "back",
    phone = nil -- phone number assigned by server
}

-- Load/save config
local function saveConfig()
    local f = fs.open(CONFIG_PATH, "w")
    if f then f.write(textutils.serialize(config)); f.close() end
end
local function loadConfig()
    if fs.exists(CONFIG_PATH) then
        local f = fs.open(CONFIG_PATH, "r")
        local data = f.readAll()
        f.close()
        local ok, t = pcall(textutils.unserialize, data)
        if ok and type(t) == "table" then
            for k,v in pairs(t) do config[k] = v end
        end
    end
end
loadConfig()

-- Ensure startup points to firm.lpsw (keeps prior behavior)
if not fs.exists(STARTUP_PATH) then
    local f = fs.open(STARTUP_PATH, "w")
    f.write('shell.run("firm.lpsw")')
    f.close()
end

-- ensure rednet open if modem present
if peripheral.find("modem") then
    rednet.open(config.rednetSide)
end

-- UTILITIES
local function centerText(y, text, color)
    local w = term.getSize()
    local x = math.floor((w - #text) / 2)
    term.setCursorPos(math.max(1, x), y)
    if color then term.setTextColor(color) end
    term.write(text)
    term.setTextColor(colors.white)
end

-- beep helper: speaker peripheral preferred, otherwise terminal bell
local function beep()
    local sp = peripheral.find("speaker")
    if sp and type(sp.playNote) == "function" then
        pcall(function() sp.playNote("harp", 5, 1) end)
    else
        io.write("\7")
    end
end

-- ===== Networking helpers =====

-- Normalize rednet.lookup output so it ALWAYS returns a table of server IDs
local function findServers()
    local ok, lookup = pcall(rednet.lookup, "lemonphone_server")
    if not ok then return {} end
    if lookup == nil then
        return {}
    end
    if type(lookup) == "number" then
        return { lookup }
    elseif type(lookup) == "table" then
        -- ensure all entries are numbers
        local out = {}
        for _, v in ipairs(lookup) do
            if type(v) == "number" then table.insert(out, v) end
        end
        return out
    else
        return {}
    end
end

-- Attempt to register with a specific server (returns ok, info)
local function registerWithServerOnce(serverID, timeout)
    timeout = timeout or 5
    local clientID = os.getComputerID()
    local payload = { type="register_request", desired_service=2, clientID = clientID, username = config.username }
    local okSend = pcall(rednet.send, serverID, payload, "lemonphone_control")
    if not okSend then return false, "send failed" end
    local sender, resp = rednet.receive("lemonphone_control", timeout)
    if sender == serverID and type(resp) == "table" and resp.type == "register_response" and resp.ok then
        return true, resp.phone
    end
    return false, "no response"
end

-- High-level ensure registration: try all discovered servers
local function ensureRegistration()
    if config.phone then return true, config.phone end
    local servers = findServers()
    if #servers == 0 then return false, "no servers found" end
    for _, sid in ipairs(servers) do
        local ok, res = pcall(registerWithServerOnce, sid, 5)
        if ok and res then
            -- registerWithServerOnce returns (true, phone) on success
            if type(res) == "string" then
                config.phone = res
                saveConfig()
                return true, res
            elseif type(res) == "table" and res.phone then
                config.phone = res.phone
                saveConfig()
                return true, config.phone
            end
        end
    end
    return false, "registration failed"
end

-- send a global chat message: we forward to one known server and let it broadcast/relay
local function sendGlobal(text)
    local servers = findServers()
    if #servers == 0 then return false, "no servers" end
    local payload = { from_phone = config.phone or ("u-" .. tostring(os.getComputerID())), text = text }
    local ok, err = pcall(rednet.send, servers[1], payload, "lemonphone_global")
    if not ok then return false, "send failed" end
    return true
end

-- send a DM (server-forward style)
local function sendDM(destPhone, text)
    local servers = findServers()
    if #servers == 0 then return false, "no servers" end
    local payload = { type = "server_forward", from_phone = config.phone, dest_phone = destPhone, text = text }
    local ok, err = pcall(rednet.send, servers[1], payload, "lemonphone_control")
    if not ok then return false, "send failed" end
    return true
end

-- Chat message storage (runtime)
local messages = {} -- array of {type="global"/"dm", from, to, text, time}
local function storeMessage(msg)
    msg.time = os.time()
    table.insert(messages, msg)
    -- cap messages to reasonable history to avoid unbounded growth
    if #messages > 500 then
        table.remove(messages, 1)
    end
end

-- ===== Background listener (runs in its own parallel thread) =====
-- listens for: lemonphone_global, lemonphone_dm, lemonphone_control
local function backgroundListener()
    while true do
        local sender, msg, proto = rednet.receive()
        if not msg then
            -- continue if receive returned nil for some reason
        else
            if proto == "lemonphone_global" then
                -- msg expected: { from_phone = "...", text = "..." }
                storeMessage({ type="global", from = msg.from_phone, text = msg.text })
                beep()
                os.queueEvent("lemon_message", msg)
            elseif proto == "lemonphone_dm" then
                -- msg expected: { from_phone = "...", to_phone = "...", text = "..." }
                storeMessage({ type="dm", from = msg.from_phone, to = msg.to_phone, text = msg.text })
                beep()
                os.queueEvent("lemon_message", msg)
            elseif proto == "lemonphone_control" then
                -- servers may forward dm_incoming as control messages
                if type(msg) == "table" and msg.type == "dm_incoming" then
                    storeMessage({ type="dm", from = msg.from_phone, to = config.phone, text = msg.text })
                    beep()
                    os.queueEvent("lemon_message", msg)
                elseif type(msg) == "table" and msg.type == "register_response" then
                    -- registration responses are generally handled elsewhere, but store for completeness
                    if msg.ok and msg.phone then
                        config.phone = msg.phone
                        saveConfig()
                    end
                end
            end
        end
    end
end

-- ===== Signal updater (estimates distance to nearest server) =====
local signalStrength = 0 -- 0..3
local function distanceToBars(distance)
    if distance <= 5 then return 3
    elseif distance <= 20 then return 2
    else return 1 end
end

local function updateSignal()
    while true do
        local servers = findServers()
        if #servers == 0 then
            signalStrength = 0
        else
            local me = os.getComputerID()
            local best = nil
            for _, s in ipairs(servers) do
                if type(s) == "number" then
                    local d = math.abs(me - s)
                    if not best or d < best then best = d end
                end
            end
            if best then signalStrength = distanceToBars(best) else signalStrength = 1 end
        end
        sleep(5)
    end
end

local function drawSignalBars()
    local barChar = "|"
    for i = 1, 3 do
        if i <= signalStrength then term.setTextColor(colors.lime) else term.setTextColor(colors.gray) end
        term.write(barChar)
    end
    term.setTextColor(colors.white)
end

-- ===== UI components =====

-- Chat UI: shows recent messages and lets user send; /dm <phone> <text> for direct messages
local function chatUI()
    while true do
        term.setBackgroundColor(config.bgColor)
        term.clear()
        centerText(2, "Messages", config.accentColor)
        term.setTextColor(colors.white)
        local start = math.max(1, #messages - 11)
        local line = 4
        for i = start, #messages do
            local m = messages[i]
            local prefix
            if m.type == "global" then
                prefix = "[" .. tostring(m.from or "?") .. "]"
            else
                local toPart = m.to and (">" .. tostring(m.to)) or ""
                prefix = "DM[" .. tostring(m.from or "?") .. toPart .. "]"
            end
            term.setCursorPos(2, line)
            term.clearLine()
            local txt = (m.text or "")
            if #prefix + 1 + #txt > term.getSize() - 4 then
                txt = txt:sub(1, term.getSize() - 4 - #prefix - 1) .. "…"
            end
            term.write(prefix .. " " .. txt)
            line = line + 1
        end

        term.setCursorPos(2, 17)
        term.setTextColor(colors.gray)
        term.clearLine()
        term.write("Type message. Commands: /dm <phone> <text>  /back")
        term.setCursorPos(2, 18)
        term.setTextColor(colors.white)
        term.clearLine()
        local input = read()
        if not input then input = "" end
        if input == "/back" then return end

        local dmDest, dmText = input:match("^/dm%s+([%w%-]+)%s+(.+)")
        if dmDest and dmText then
            local ok, err = sendDM(dmDest, dmText)
            if ok then
                storeMessage({ type="dm", from = config.phone, to = dmDest, text = dmText })
            else
                -- show brief error
                term.setCursorPos(2, 18)
                term.setTextColor(colors.red)
                term.clearLine()
                term.write("Failed to send DM: " .. tostring(err))
                sleep(1.5)
            end
        else
            if input ~= "" then
                local ok, err = sendGlobal(input)
                if ok then
                    storeMessage({ type="global", from = config.phone, text = input })
                else
                    term.setCursorPos(2, 18)
                    term.setTextColor(colors.red)
                    term.clearLine()
                    term.write("Failed to send: " .. tostring(err))
                    sleep(1.5)
                end
            end
        end
    end
end

-- About page
local function aboutLemonphone()
    term.setBackgroundColor(config.bgColor)
    term.clear()
    centerText(3,"About Lemonphone", config.accentColor)
    term.setCursorPos(2,5)
    term.setTextColor(colors.white)
    term.write("Device: " .. firmware.device)
    term.setCursorPos(2,6)
    term.write("Firmware: " .. firmware.version)
    term.setCursorPos(2,7)
    term.write("Phone: " .. tostring(config.phone or "unregistered"))
    term.setCursorPos(2,9)
    term.setTextColor(colors.gray)
    term.write("Press any key to return...")
    os.pullEvent("key")
end

-- Settings UI
local function openSettings()
    while true do
        term.setBackgroundColor(config.bgColor)
        term.clear()
        centerText(3, "Settings", config.accentColor)
        term.setCursorPos(2,5)
        term.setTextColor(colors.white)
        term.write("1. Change Username")
        term.setCursorPos(2,6)
        term.write("2. Change Accent Color")
        term.setCursorPos(2,7)
        term.write("3. Re-register Phone")
        term.setCursorPos(2,8)
        term.write("4. Back")
        local _, key = os.pullEvent("key")
        if key == keys.one then
            term.setCursorPos(2,10); term.clearLine(); term.setTextColor(colors.white)
            term.write("Enter new username: ")
            local newName = read()
            if newName and newName ~= "" then config.username = newName; saveConfig() end
        elseif key == keys.two then
            term.setCursorPos(2,10); term.clearLine(); term.setTextColor(colors.white)
            term.write("Color name (e.g. yellow): ")
            local colorName = read()
            if colorName and colors[colorName] then config.accentColor = colors[colorName]; saveConfig()
            else term.setCursorPos(2,11); term.setTextColor(colors.red); term.write("Invalid color."); sleep(1) end
        elseif key == keys.three then
            -- force re-register (remove phone then try)
            config.phone = nil; saveConfig()
            term.setCursorPos(2,10); term.clearLine(); term.setTextColor(colors.white)
            term.write("Attempting registration...")
            local ok, res = ensureRegistration()
            term.setCursorPos(2,11); term.clearLine()
            if ok then term.setTextColor(colors.lime); term.write("Registered as " .. res)
            else term.setTextColor(colors.red); term.write("Registration failed: " .. tostring(res)) end
            sleep(1.5)
        elseif key == keys.four then
            return
        end
    end
end

-- Top bar and menu drawing
local function drawTopBar()
    local w,_ = term.getSize()
    term.setCursorPos(1,1)
    term.setBackgroundColor(config.accentColor)
    term.clearLine()
    term.setCursorPos(2,1)
    if rednet.isOpen(config.rednetSide) then
        drawSignalBars()
        term.write(" Net")
    else
        term.setTextColor(colors.red)
        term.write("X NoNet")
        term.setTextColor(colors.white)
    end
    local timeStr = textutils.formatTime(os.time(), true)
    local dateStr = os.date("%b %d")
    local clock = timeStr .. " | " .. dateStr
    term.setCursorPos(w - #clock - 1, 1)
    term.setTextColor(colors.white)
    term.write(clock)
end

local function drawMenu(selected)
    local apps = {
        { name="Messages", fn = chatUI },
        { name="Files", fn = function() shell.run("shell") end },
        { name="Worm", fn = function() pcall(shell.run, "worm") end },
        { name="Falling Blocks", fn = function() pcall(shell.run, "falling") end },
        { name="Settings", fn = openSettings },
        { name="About Lemonphone", fn = aboutLemonphone },
        { name="Shutdown", fn = function() os.shutdown() end },
    }

    local startY = 4
    term.setTextColor(colors.white)
    for i, app in ipairs(apps) do
        term.setCursorPos(4, startY + i - 1)
        if i == selected then term.setBackgroundColor(colors.gray); term.setTextColor(colors.black)
        else term.setBackgroundColor(config.bgColor); term.setTextColor(colors.white) end
        term.clearLine()
        term.setCursorPos(4, startY + i - 1)
        term.write(app.name)
    end
    term.setBackgroundColor(config.bgColor)
    return apps
end

-- Main UI loop
local function mainLoop()
    -- Try auto-register but don't block UI long
    spawn(function() ensureRegistration() end)

    local selected = 1
    local refreshRate = 1
    while true do
        term.setBackgroundColor(config.bgColor); term.clear()
        drawTopBar()
        centerText(3, (config.username or "User") .. "'s Lemonphone", config.accentColor)
        local apps = drawMenu(selected)
        local timerId = os.startTimer(refreshRate)
        local event, a, b, c = os.pullEvent()
        if event == "key" then
            if a == keys.up then selected = math.max(1, selected - 1)
            elseif a == keys.down then selected = math.min(#apps, selected + 1)
            elseif a == keys.enter then
                local chosen = apps[selected]
                term.setBackgroundColor(config.bgColor); term.clear(); term.setCursorPos(1,1)
                pcall(chosen.fn)
            end
        elseif event == "timer" and a == timerId then
            -- redraw topbar only
            drawTopBar()
        elseif event == "lemon_message" then
            -- show short banner
            local msg = a -- the queued message
            term.setCursorPos(2,2)
            term.setTextColor(colors.white)
            term.clearLine()
            term.write("[New message]")
            sleep(1)
        end
    end
end

-- Small spawn wrapper (simple co-routine starter)
function spawn(f)
    local co = coroutine.create(f)
    local function step()
        local ok, err = coroutine.resume(co)
        if not ok and err then print("spawn error: "..tostring(err)) end
    end
    -- run in background via a timer loop
    parallel.waitForAny(function() while coroutine.status(co) ~= "dead" do step(); sleep(0.05) end end, function() end)
end

-- But the above spawn uses parallel.waitForAny blocking — replace spawn with simpler nonblocking approach below
-- We'll instead start background tasks with parallel.waitForAny at script bottom.

-- STARTUP: run background tasks in parallel
local function startAll()
    -- Use three parallel tasks:
    -- 1) signal updater
    -- 2) background listener
    -- 3) mainLoop
    parallel.waitForAny(updateSignal, backgroundListener, mainLoop)
end

-- Start
term.setBackgroundColor(colors.black)
term.clear()
print("Starting LemonOS 3.0 (fixed)...")
startAll()