ARK-OS/arkos/arkrt/compositor/ArkGraphicsServer.swift

//
// Copyright 2026 Aarav Ravindra Kharade
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

// ═══════════════════════════════════════════════════════════════════
//  ARK-Graphics Server (Display + Registry)
// ═══════════════════════════════════════════════════════════════════
//  Implements wl_display (object ID 1) and wl_registry.
//  Manages global interface registration and client lifecycle.
// ═══════════════════════════════════════════════════════════════════

import ark.sys.libc

// MARK: - Global Interface

/// A registered global interface that clients can bind to.
public struct WaylandGlobal {
    public let name: UInt32           // Unique numeric name for this global
    public let interface: String       // Interface name (e.g., "wl_compositor")
    public let version: UInt32         // Interface version
    public let factory: (ArkClient, UInt32, UInt32) -> WaylandObject?  // (client, id, version) → object

    public init(name: UInt32, interface: String, version: UInt32,
                factory: @escaping (ArkClient, UInt32, UInt32) -> WaylandObject?) {
        self.name = name
        self.interface = interface
        self.version = version
        self.factory = factory
    }
}

// MARK: - WlDisplay Object

/// The wl_display server-side object. Always object ID 1 for every client.
public class WlDisplay: WaylandObject {
    public let interfaceName = "wl_display"
    public weak var server: ArkGraphicsServer?

    public init(server: ArkGraphicsServer) {
        self.server = server
    }

    public func handleRequest(client: ArkClient, opcode: UInt16, decoder: inout WaylandDecoder) -> Bool {
        guard let req = WlDisplayRequest(rawValue: opcode) else { return false }

        switch req {
        case .sync:
            // Client sends: sync(callback: new_id<wl_callback>)
            guard let callbackID = decoder.readNewID() else { return false }
            let callback = WlCallback()
            _ = client.registerObject(id: callbackID, object: callback)

            // Immediately fire the callback with current serial
            var encoder = WaylandEncoder()
            encoder.writeUInt32(server?.nextSerial() ?? 0)
            let event = encoder.buildMessage(objectID: callbackID, opcode: WlCallbackEvent.done.rawValue)
            client.queueEvent(event)

            // The callback is now consumed — destroy it
            client.destroyObject(id: callbackID)
            return true

        case .getRegistry:
            // Client sends: get_registry(registry: new_id<wl_registry>)
            guard let registryID = decoder.readNewID() else { return false }
            guard let server = server else { return false }

            let registry = WlRegistry(server: server)
            _ = client.registerObject(id: registryID, object: registry)

            // Send all current globals to the new registry
            for global in server.globals {
                var encoder = WaylandEncoder()
                encoder.writeUInt32(global.name)
                encoder.writeString(global.interface)
                encoder.writeUInt32(global.version)
                let event = encoder.buildMessage(objectID: registryID, opcode: WlRegistryEvent.global.rawValue)
                client.queueEvent(event)
            }
            return true
        }
    }

    public func destroy(client: ArkClient) {
        // wl_display is never truly destroyed
    }
}

// MARK: - WlCallback Object

/// A one-shot callback object (used by wl_display.sync and wl_surface.frame).
public class WlCallback: WaylandObject {
    public let interfaceName = "wl_callback"

    public func handleRequest(client: ArkClient, opcode: UInt16, decoder: inout WaylandDecoder) -> Bool {
        // wl_callback has no requests
        return false
    }
}

// MARK: - WlRegistry Object

/// The wl_registry object — handles global binding.
public class WlRegistry: WaylandObject {
    public let interfaceName = "wl_registry"
    public weak var server: ArkGraphicsServer?

    public init(server: ArkGraphicsServer) {
        self.server = server
    }

    public func handleRequest(client: ArkClient, opcode: UInt16, decoder: inout WaylandDecoder) -> Bool {
        guard let req = WlRegistryRequest(rawValue: opcode) else { return false }

        switch req {
        case .bind:
            // Client sends: bind(name: uint, interface: string, version: uint, id: new_id)
            // Note: for wl_registry.bind, the new_id arg includes interface+version
            guard let name = decoder.readUInt32(),
                  let interface = decoder.readString(),
                  let version = decoder.readUInt32(),
                  let newID = decoder.readNewID(),
                  let server = server else { return false }

            // Find the global
            guard let global = server.globals.first(where: { $0.name == name }) else {
                client.sendError(objectID: 1, code: .invalidObject,
                                message: "No global with name \(name)")
                return true
            }

            // Verify interface matches
            guard global.interface == interface else {
                client.sendError(objectID: 1, code: .invalidObject,
                                message: "Global \(name) is \(global.interface), not \(interface)")
                return true
            }

            // Create the bound object
            let boundVersion = min(version, global.version)
            guard let object = global.factory(client, newID, boundVersion) else {
                client.sendError(objectID: 1, code: .noMemory,
                                message: "Failed to create \(interface)")
                return true
            }

            _ = client.registerObject(id: newID, object: object)
            return true
        }
    }
}

// MARK: - ArkGraphicsServer

/// The top-level display server. Manages globals, clients, and the event loop.
public class ArkGraphicsServer {
    public let socket: ArkSocket
    public var clients: [Int32: ArkClient] = [:]
    public var globals: [WaylandGlobal] = []
    public var running: Bool = true
    
    public weak var compositor: ArkCompositor?

    /// Monotonically increasing serial number for events.
    private var serial: UInt32 = 0

    /// Next global name (auto-increment).
    private var nextGlobalName: UInt32 = 1

    public init?(socketPath: String = "/run/ark-graphics.sock") {
        guard let sock = ArkSocket(path: socketPath) else {
            return nil
        }
        self.socket = sock
    }

    /// Allocates the next event serial number.
    public func nextSerial() -> UInt32 {
        serial += 1
        return serial
    }

    // MARK: - Global Registration

    /// Registers a new global interface.
    public func addGlobal(interface: String, version: UInt32,
                          factory: @escaping (ArkClient, UInt32, UInt32) -> WaylandObject?) {
        let name = nextGlobalName
        nextGlobalName += 1

        let global = WaylandGlobal(name: name, interface: interface,
                                   version: version, factory: factory)
        globals.append(global)

        // Notify all existing clients' registries
        for (_, client) in clients {
            for (id, obj) in client.objects {
                if obj is WlRegistry {
                    var encoder = WaylandEncoder()
                    encoder.writeUInt32(name)
                    encoder.writeString(interface)
                    encoder.writeUInt32(version)
                    let event = encoder.buildMessage(objectID: id, opcode: WlRegistryEvent.global.rawValue)
                    client.queueEvent(event)
                }
            }
        }
    }

    // MARK: - Client Management

    /// Creates a new client from an accepted connection.
    public func addClient(fd: Int32) {
        let client = ArkClient(fd: fd, server: self)

        // Pre-register wl_display as object ID 1
        let display = WlDisplay(server: self)
        _ = client.registerObject(id: 1, object: display)

        clients[fd] = client
        print("ARK-Graphics: Client connected (fd=\(fd))")
    }

    /// Removes and cleans up a disconnected client.
    public func removeClient(fd: Int32) {
        if let client = clients.removeValue(forKey: fd) {
            // Destroy all client objects
            for (_, obj) in client.objects {
                obj.destroy(client: client)
            }
            client.objects.removeAll()
            socket.removeFromEpoll(fd: fd)
            close(fd)
            print("ARK-Graphics: Client disconnected (fd=\(fd))")
        }
    }

    // MARK: - Event Loop

    /// Runs a single iteration of the event loop.
    /// Returns the events that were processed (for compositor integration).
    public func pollOnce(timeout: Int32 = 16) {
        let events = socket.wait(timeout: timeout)

        for (fd, eventMask) in events {
            if fd == socket.listenFD {
                // New client connection
                if let clientFD = socket.acceptClient() {
                    addClient(fd: clientFD)
                }
            } else if let client = clients[fd] {
                if eventMask & (EPOLLERR | EPOLLHUP) != 0 {
                    // Client error or hangup
                    removeClient(fd: fd)
                } else if eventMask & EPOLLIN != 0 {
                    // Incoming data from client
                    if let (bytes, fds) = socket.recvMessage(from: fd) {
                        client.recvBuffer.append(contentsOf: bytes)
                        client.recvFDs.append(contentsOf: fds)
                        client.dispatchMessages()
                    } else {
                        // recvMessage returned nil — client disconnected
                        removeClient(fd: fd)
                    }
                }
            }
        }

        // Flush all client event queues
        for (_, client) in clients {
            if client.alive {
                client.flush(socket: socket)
            } else {
                removeClient(fd: client.fd)
            }
        }
    }
}