-- lemonos3.lua
-- LemonOS 3.0 - integrated chat + DM + multi-server relay client
-- Put as /lemonos3.lua and have startup run it (or replace existing lemonos)

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

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

-- load/save
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()

-- startup ensure
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 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(x, y)
    if color then term.setTextColor(color) end
    term.write(text)
    term.setTextColor(colors.white)
end

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

-- === PHONE REGISTRATION ===
-- find servers
local function findServers()
    local servers = rednet.lookup("lemonphone_server") or {}
    return servers
end

local function registerWithServer(serverID)
    local clientID = os.getComputerID()
    rednet.send(serverID, { type="register_request", desired_service=2, clientID = clientID }, "lemonphone_control")
    local sender, resp, proto = rednet.receive("lemonphone_control", 5)
    if sender == serverID and resp and resp.type == "register_response" and resp.ok then
        config.phone = resp.phone
        saveConfig()
        return true, resp.phone
    end
    return false, "no response"
end

-- attempt auto-registration (non-blocking friendly)
local function ensureRegistration()
    if config.phone then return true, config.phone end
    local servers = findServers()
    if 1 == 0 then return false, "no servers" end --Temp fix for crash, actual resolution in lemonOS 3.0.1
    for _,sid in ipairs(servers) do
        local ok, phoneOrErr = pcall(registerWithServer, sid)
        if ok and phoneOrErr then
            return true, phoneOrErr
        end
    end
    return false, "failed"
end

-- === CHAT STORAGE ===
local messages = {} -- list of {type="global"|"dm", from_phone, to_phone, text, time, seen}
local function storeMessage(m)
    m.time = os.time()
    table.insert(messages, m)
end

-- === BACKGROUND LISTENER ===
-- listens for: lemonphone_global, lemonphone_dm, lemonphone_control (server responses)
local function backgroundListener()
    while true do
        local sender, msg, proto = rednet.receive()
        if proto == "lemonphone_global" then
            -- broadcast to local clients; store message for UI
            storeMessage({ type="global", from_phone = msg.from_phone, text = msg.text })
            -- beep and notify
            beep()
            os.queueEvent("lemon_message", msg)
        elseif proto == "lemonphone_dm" then
            -- direct DM to this client
            storeMessage({ type="dm", from_phone = msg.from_phone, to_phone = config.phone, text = msg.text })
            beep()
            os.queueEvent("lemon_message", msg)
        elseif proto == "lemonphone_control" then
            -- server replies or forwarded DM from server
            if msg.type == "dm_incoming" then
                storeMessage({ type="dm", from_phone = msg.from_phone, to_phone = config.phone, text = msg.text })
                beep()
                os.queueEvent("lemon_message", msg)
            elseif msg.type == "register_response" then
                -- handled by registerWithServer; we still store
            end
        end
    end
end

-- === SENDER HELPERS ===
-- send global message (to server which will broadcast/relay)
local function sendGlobal(text)
    -- choose a server (first available)
    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 }
    -- send with protocol lemonphone_global to server then server handles broadcast
    -- servers also respond by rebroadcast so we just send to any server
    rednet.send(servers[1], payload, "lemonphone_global")
    return true
end

-- send DM via appropriate server: server will forward to owner server
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 }
    -- send to the first server; servers will forward to correct owner
    rednet.send(servers[1], payload, "lemonphone_control")
    return true
end

-- === UI: Chat window ===
local function chatUI()
    -- very simple: show last 12 messages, allow typing, send to global or DM when prefixed with /dm <phone> <text>
    while true do
        term.setBackgroundColor(config.bgColor)
        term.clear()
        centerText(2, "Messages", config.accentColor)
        term.setCursorPos(2,4)
        term.setTextColor(colors.white)
        local start = math.max(1, #messages - 11)
        for i = start, #messages do
            local m = messages[i]
            local prefix = (m.type == "global") and ("["..tostring(m.from_phone).."]") or ("DM ["..tostring(m.from_phone)..">"..tostring(m.to_phone).."]")
            term.setCursorPos(2, 4 + (i - start))
            term.clearLine()
            term.write(prefix .. " " .. (m.text or ""))
        end
        term.setCursorPos(2, 17)
        term.clearLine()
        term.setTextColor(colors.gray)
        term.write("Type message. Commands: /dm <phone> <text>  /back")
        term.setCursorPos(2, 18)
        term.setTextColor(colors.white)
        local input = read()
        if input == "/back" then return end
        local dmMatch = string.match(input, "^/dm%s+([%w%-]+)%s+(.+)")
        if dmMatch then
            local dest, msgText = string.match(input, "^/dm%s+([%w%-]+)%s+(.+)")
            if dest and msgText then
                sendDM(dest, msgText)
                storeMessage({ type="dm", from_phone = config.phone, to_phone = dest, text = msgText })
            end
        else
            if input ~= "" then
                sendGlobal(input)
                storeMessage({ type="global", from_phone = config.phone, text = input })
            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.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.write("Press any key to return...")
    os.pullEvent("key")
end

-- === Settings ===
local function openSettings()
    while true do
        term.setBackgroundColor(config.bgColor)
        term.clear()
        centerText(3, "Settings", config.accentColor)
        term.setCursorPos(2,5)
        term.write("1. Change Username")
        term.setCursorPos(2,6)
        term.write("2. Change Accent Color")
        term.setCursorPos(2,7)
        term.write("3. Back")
        local _, key = os.pullEvent("key")
        if key == keys.one then
            term.setCursorPos(2,9); term.clearLine(); term.write("Enter new username: ")
            config.username = read(); saveConfig()
        elseif key == keys.two then
            term.setCursorPos(2,9); term.clearLine(); term.write("Color name (e.g. yellow): ")
            local colorName = read()
            if colors[colorName] then config.accentColor = colors[colorName]; saveConfig()
            else term.setCursorPos(2,10); term.write("Invalid color."); sleep(1) end
        elseif key == keys.three then return end
    end
end

-- === Top bar and signal (non-blocking) ===
local signalStrength = 0
local function distanceToBars(distance)
    if distance <= 5 then return 3
    elseif distance <= 20 then return 2
    else return 1 end
end

-- async updater: find servers and estimate distance to nearest
local function updateSignal()
    while true do
        local servers = findServers()
        if #servers == 0 then signalStrength = 0
        else
            local me = os.getComputerID()
            local bestd = nil
            for _,s in ipairs(servers) do
                local d = math.abs(me - s)
                if not bestd or d < bestd then bestd = d end
            end
            if bestd then signalStrength = distanceToBars(bestd) 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

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

-- === Main UI ===
local function drawMenu(selected)
    local apps = {
        {name="Messages", fn=chatUI},
        {name="Files", fn=function() shell.run("shell") end},
        {name="Worm", fn=function() shell.run("worm") end},
        {name="Falling Blocks", fn=function() 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 loop
local function mainLoop()
    -- ensure registration attempt in background
    local ok, phone = ensureRegistration()
    if ok then
        -- if server accepted, nothing else
    end

    local selected = 1
    local refreshRate = 1
    os.startTimer(0) -- kick timers
    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 timer = os.startTimer(refreshRate)
        local event, a = 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)
                chosen.fn()
            end
        elseif event == "timer" then
            -- redraw topbar only
            drawTopBar()
        elseif event == "lemon_message" then
            -- optional: show small banner at top for a second
            term.setCursorPos(2,2)
            term.setTextColor(colors.white)
            term.write("[New message]")
            sleep(1)
        end
    end
end

-- === Startup: background tasks ===
local function startAll()
    parallel.waitForAny(updateSignal, backgroundListener, mainLoop)
end

-- ensure rednet host presence? not needed for client
print("Starting LemonOS 3.0...")
startAll()