/* SPDX-License-Identifier: Apache-2.0 */
/*
 * Schcom -- the SchOS compositor.
 *
 * Owns the display and the input devices. Clients draw into shared buffers and
 * hand them over; Schcom decides what is on top of what, composites, and
 * presents. It does not decide where windows go: that is SchWM's job, and the
 * moment this file starts making placement decisions the split has failed.
 *
 * v0, deliberately small but not a toy:
 *
 *   - real framebuffer output through /dev/fb0 and the fbdev ioctls
 *   - real clients over a Unix socket, with a versioned handshake
 *   - real shared-memory buffers, one per surface
 *   - a back buffer, so a half-composited frame never reaches the screen
 *   - input read from /dev/input/eventN and routed to the top surface
 *
 * What it does not do yet: damage tracking beyond "repaint everything",
 * multiple outputs, subsurfaces, or any notion of focus beyond "topmost".
 * Those are additions, not rewrites -- which is the point of committing to the
 * surface/commit model now rather than growing into it.
 */

#define _GNU_SOURCE
#include <errno.h>
#include <fcntl.h>
#include <poll.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/ioctl.h>
#include <sys/mman.h>
#include <sys/shm.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <sys/un.h>
#include <linux/fb.h>
#include <linux/input.h>

#include "proto.h"

#define MAX_CLIENTS  16
#define MAX_SURFACES 32

struct surface {
    int      used;
    uint32_t id;
    int      client;        /* index into clients[] */
    int      shmid;
    void    *pixels;
    uint32_t w, h, stride, size;
    int32_t  x, y;
};

struct client {
    int  fd;                /* -1 when the slot is free */
    int  hello_done;
};

static struct client  clients[MAX_CLIENTS];
static struct surface surfaces[MAX_SURFACES];
static uint32_t       next_surface_id = 1;

/* The screen. */
static uint8_t *fb;
static uint32_t fb_w, fb_h, fb_stride, fb_bpp;
static size_t   fb_size;
static uint8_t *back;           /* composited here, then copied out in one go */

/* Pointer position, moved by EV_REL from the mouse. */
static int32_t ptr_x, ptr_y;

static void
die(const char *what)
{
    fprintf(stderr, "schcom: %s: %s\n", what, strerror(errno));
    exit(1);
}

/* ---------------------------------------------------------------- output */

static int
open_framebuffer(void)
{
    int fd = open("/dev/fb0", O_RDWR);
    if (fd < 0)
        die("open /dev/fb0");

    struct fb_var_screeninfo var;
    struct fb_fix_screeninfo fix;
    if (ioctl(fd, FBIOGET_VSCREENINFO, &var) < 0)
        die("FBIOGET_VSCREENINFO");
    if (ioctl(fd, FBIOGET_FSCREENINFO, &fix) < 0)
        die("FBIOGET_FSCREENINFO");

    fb_w      = var.xres;
    fb_h      = var.yres;
    fb_bpp    = var.bits_per_pixel;
    /* The stride comes from the driver, never from xres*bpp: the two differ on
     * real hardware and assuming otherwise shears the image. */
    fb_stride = fix.line_length;
    fb_size   = (size_t)fb_stride * fb_h;

    if (fb_bpp != 32) {
        fprintf(stderr, "schcom: need 32bpp, got %u\n", fb_bpp);
        exit(1);
    }

    fb = mmap(NULL, fb_size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
    if (fb == MAP_FAILED)
        die("mmap /dev/fb0");

    back = malloc(fb_size);
    if (!back)
        die("malloc back buffer");

    printf("schcom: %ux%u, %u bpp, stride %u\n", fb_w, fb_h, fb_bpp, fb_stride);
    fflush(stdout);
    return fd;
}

static void
composite(void)
{
    /* Desktop. A flat colour rather than black, so that "nothing was drawn"
     * and "the compositor is not running" cannot look the same. */
    for (uint32_t y = 0; y < fb_h; y++) {
        uint32_t *row = (uint32_t *)(back + (size_t)y * fb_stride);
        for (uint32_t x = 0; x < fb_w; x++)
            row[x] = 0x00203040;
    }

    /* Surfaces, bottom to top. Creation order is stacking order in v0; SchWM
     * will reorder later through a protocol request rather than by reaching in
     * here. */
    for (int i = 0; i < MAX_SURFACES; i++) {
        struct surface *s = &surfaces[i];
        if (!s->used || !s->pixels)
            continue;

        for (uint32_t sy = 0; sy < s->h; sy++) {
            int32_t dy = s->y + (int32_t)sy;
            if (dy < 0 || dy >= (int32_t)fb_h)
                continue;

            uint32_t *dst = (uint32_t *)(back + (size_t)dy * fb_stride);
            uint32_t *src = (uint32_t *)((uint8_t *)s->pixels + (size_t)sy * s->stride);

            for (uint32_t sx = 0; sx < s->w; sx++) {
                int32_t dx = s->x + (int32_t)sx;
                if (dx < 0 || dx >= (int32_t)fb_w)
                    continue;
                dst[dx] = src[sx];
            }
        }
    }

    /* A crude pointer, drawn last so it is never under a window. */
    for (int32_t py = 0; py < 10; py++) {
        int32_t dy = ptr_y + py;
        if (dy < 0 || dy >= (int32_t)fb_h)
            continue;
        uint32_t *dst = (uint32_t *)(back + (size_t)dy * fb_stride);
        for (int32_t px = 0; px < 10 - py; px++) {
            int32_t dx = ptr_x + px;
            if (dx >= 0 && dx < (int32_t)fb_w)
                dst[dx] = 0x00FFFFFF;
        }
    }

    /* One copy to the screen. Compositing straight into the framebuffer would
     * show every intermediate state -- the flicker would be the whole desktop
     * being repainted under the user's eyes. */
    memcpy(fb, back, fb_size);
}

/* ---------------------------------------------------------------- clients */

static int
send_all(int fd, const void *buf, size_t n)
{
    const uint8_t *p = buf;
    while (n) {
        ssize_t w = write(fd, p, n);
        if (w < 0) {
            if (errno == EINTR)
                continue;
            return -1;
        }
        p += w;
        n -= (size_t)w;
    }
    return 0;
}

static int
read_all(int fd, void *buf, size_t n)
{
    uint8_t *p = buf;
    while (n) {
        ssize_t r = read(fd, p, n);
        if (r == 0)
            return -1;              /* client went away */
        if (r < 0) {
            if (errno == EINTR)
                continue;
            return -1;
        }
        p += r;
        n -= (size_t)r;
    }
    return 0;
}

static void
send_error(int fd, uint32_t serial, uint32_t code)
{
    struct schcom_error e;
    memset(&e, 0, sizeof e);
    e.hdr.op     = SCHCOM_OP_ERROR;
    e.hdr.length = sizeof e;
    e.hdr.serial = serial;
    e.code       = code;
    send_all(fd, &e, sizeof e);
}

static void
drop_client(int ci)
{
    if (clients[ci].fd < 0)
        return;

    for (int i = 0; i < MAX_SURFACES; i++) {
        struct surface *s = &surfaces[i];
        if (s->used && s->client == ci) {
            if (s->pixels)
                shmdt(s->pixels);
            /* IPC_RMID once, here: the segment goes when the last attach does,
             * so a client that died mid-frame does not leak it. */
            shmctl(s->shmid, IPC_RMID, NULL);
            memset(s, 0, sizeof *s);
        }
    }

    close(clients[ci].fd);
    clients[ci].fd = -1;
    clients[ci].hello_done = 0;
    printf("schcom: client %d gone\n", ci);
    fflush(stdout);
}

static void
handle_create_surface(int ci, const struct schcom_create_surface *req)
{
    int fd = clients[ci].fd;

    int slot = -1;
    for (int i = 0; i < MAX_SURFACES; i++)
        if (!surfaces[i].used) { slot = i; break; }
    if (slot < 0) {
        send_error(fd, req->hdr.serial, ENOSPC);
        return;
    }

    uint32_t w = req->width, h = req->height;
    if (w == 0 || h == 0 || w > fb_w || h > fb_h) {
        send_error(fd, req->hdr.serial, EINVAL);
        return;
    }

    uint32_t stride = w * 4;
    uint32_t size   = stride * h;

    /* The compositor owns the segment, not the client: it has to survive the
     * client crashing mid-frame, and it must not be something a client can
     * quietly resize underneath us. */
    int shmid = shmget(IPC_PRIVATE, size, IPC_CREAT | 0600);
    if (shmid < 0) {
        send_error(fd, req->hdr.serial, errno);
        return;
    }
    void *pix = shmat(shmid, NULL, 0);
    if (pix == (void *)-1) {
        shmctl(shmid, IPC_RMID, NULL);
        send_error(fd, req->hdr.serial, errno);
        return;
    }
    memset(pix, 0, size);

    struct surface *s = &surfaces[slot];
    s->used   = 1;
    s->id     = next_surface_id++;
    s->client = ci;
    s->shmid  = shmid;
    s->pixels = pix;
    s->w      = w;
    s->h      = h;
    s->stride = stride;
    s->size   = size;
    s->x      = req->x;
    s->y      = req->y;

    struct schcom_surface rep;
    memset(&rep, 0, sizeof rep);
    rep.hdr.op      = SCHCOM_OP_SURFACE;
    rep.hdr.length  = sizeof rep;
    rep.hdr.serial  = req->hdr.serial;
    rep.hdr.surface = s->id;
    rep.handle_kind = SCHCOM_HANDLE_SHMID;
    rep.handle      = (uint32_t)shmid;
    rep.width       = w;
    rep.height      = h;
    rep.stride      = stride;
    rep.size        = size;

    if (send_all(fd, &rep, sizeof rep) < 0) {
        drop_client(ci);
        return;
    }
    printf("schcom: surface %u %ux%u at %d,%d for client %d (shm %d)\n",
           s->id, w, h, s->x, s->y, ci, shmid);
    fflush(stdout);
}

static struct surface *
find_surface(uint32_t id)
{
    for (int i = 0; i < MAX_SURFACES; i++)
        if (surfaces[i].used && surfaces[i].id == id)
            return &surfaces[i];
    return NULL;
}

/* Returns 1 if the screen needs repainting. */
static int
handle_message(int ci)
{
    int fd = clients[ci].fd;
    struct schcom_header hdr;

    if (read_all(fd, &hdr, sizeof hdr) < 0) {
        drop_client(ci);
        return 0;
    }
    if (hdr.length < sizeof hdr || hdr.length > 4096) {
        send_error(fd, hdr.serial, EPROTO);
        drop_client(ci);
        return 0;
    }

    uint8_t body[4096];
    memcpy(body, &hdr, sizeof hdr);
    size_t rest = hdr.length - sizeof hdr;
    if (rest && read_all(fd, body + sizeof hdr, rest) < 0) {
        drop_client(ci);
        return 0;
    }

    if (!clients[ci].hello_done && hdr.op != SCHCOM_OP_HELLO) {
        /* Nothing before the handshake. A client that gets this wrong is a
         * client built against a protocol we do not speak. */
        send_error(fd, hdr.serial, EPROTO);
        drop_client(ci);
        return 0;
    }

    switch (hdr.op) {
    case SCHCOM_OP_HELLO: {
        const struct schcom_hello *h = (const void *)body;
        if (hdr.length < sizeof *h || h->version != SCHCOM_VERSION) {
            send_error(fd, hdr.serial, EPROTONOSUPPORT);
            drop_client(ci);
            return 0;
        }
        struct schcom_hello ack;
        memset(&ack, 0, sizeof ack);
        ack.hdr.op     = SCHCOM_OP_HELLO_ACK;
        ack.hdr.length = sizeof ack;
        ack.hdr.serial = hdr.serial;
        ack.version    = SCHCOM_VERSION;
        if (send_all(fd, &ack, sizeof ack) < 0) {
            drop_client(ci);
            return 0;
        }
        clients[ci].hello_done = 1;
        printf("schcom: client %d speaks v%u\n", ci, SCHCOM_VERSION);
        fflush(stdout);
        return 0;
    }

    case SCHCOM_OP_CREATE_SURFACE:
        if (hdr.length < sizeof(struct schcom_create_surface)) {
            send_error(fd, hdr.serial, EINVAL);
            return 0;
        }
        handle_create_surface(ci, (const void *)body);
        return 1;

    case SCHCOM_OP_COMMIT: {
        struct surface *s = find_surface(hdr.surface);
        if (!s || s->client != ci) {
            send_error(fd, hdr.serial, ENOENT);
            return 0;
        }
        /* The commit is the whole point: the pixels were already there, but
         * they were not allowed to be shown until now. */
        struct schcom_header fr;
        memset(&fr, 0, sizeof fr);
        fr.op      = SCHCOM_OP_FRAME;
        fr.length  = sizeof fr;
        fr.surface = s->id;
        if (send_all(fd, &fr, sizeof fr) < 0) {
            drop_client(ci);
            return 1;
        }
        return 1;
    }

    case SCHCOM_OP_DESTROY: {
        struct surface *s = find_surface(hdr.surface);
        if (s && s->client == ci) {
            if (s->pixels)
                shmdt(s->pixels);
            shmctl(s->shmid, IPC_RMID, NULL);
            memset(s, 0, sizeof *s);
            return 1;
        }
        return 0;
    }

    default:
        /* Unknown but well-formed: skipped, not fatal. That is what the length
         * field is for, and it is how a v0 compositor survives a v1 client. */
        return 0;
    }
}

/* ----------------------------------------------------------------- input */

/* Returns 1 if the screen needs repainting, -1 to quit. */
static int
handle_input(int fd)
{
    struct input_event ev;
    if (read(fd, &ev, sizeof ev) != (ssize_t)sizeof ev)
        return 0;

    int repaint = 0;

    if (ev.type == EV_REL) {
        if (ev.code == REL_X) { ptr_x += ev.value; repaint = 1; }
        if (ev.code == REL_Y) { ptr_y += ev.value; repaint = 1; }
        if (ptr_x < 0) ptr_x = 0;
        if (ptr_y < 0) ptr_y = 0;
        if (ptr_x >= (int32_t)fb_w) ptr_x = (int32_t)fb_w - 1;
        if (ptr_y >= (int32_t)fb_h) ptr_y = (int32_t)fb_h - 1;
    }

    if (ev.type == EV_KEY && ev.value == 1 && ev.code == KEY_ESC)
        return -1;

    /* Route to the topmost surface. "Topmost" is the whole focus model in v0;
     * anything cleverer is policy and belongs to SchWM. */
    struct surface *top = NULL;
    for (int i = 0; i < MAX_SURFACES; i++)
        if (surfaces[i].used)
            top = &surfaces[i];

    if (top && (ev.type == EV_KEY || ev.type == EV_REL)) {
        struct schcom_input msg;
        memset(&msg, 0, sizeof msg);
        msg.hdr.op      = SCHCOM_OP_INPUT;
        msg.hdr.length  = sizeof msg;
        msg.hdr.surface = top->id;
        msg.type        = ev.type;
        msg.code        = ev.code;
        msg.value       = ev.value;
        msg.x           = ptr_x;
        msg.y           = ptr_y;
        if (clients[top->client].fd >= 0)
            send_all(clients[top->client].fd, &msg, sizeof msg);
    }

    return repaint;
}

/* ------------------------------------------------------------------ main */

int
main(int argc, char **argv)
{
    int frames_then_exit = 0;      /* --frames N: for automated testing */

    for (int i = 1; i < argc; i++)
        if (!strcmp(argv[i], "--frames") && i + 1 < argc)
            frames_then_exit = atoi(argv[++i]);

    for (int i = 0; i < MAX_CLIENTS; i++)
        clients[i].fd = -1;

    open_framebuffer();
    ptr_x = (int32_t)fb_w / 2;
    ptr_y = (int32_t)fb_h / 2;

    unlink(SCHCOM_SOCKET_PATH);
    int lfd = socket(AF_UNIX, SOCK_STREAM, 0);
    if (lfd < 0)
        die("socket");

    struct sockaddr_un addr;
    memset(&addr, 0, sizeof addr);
    addr.sun_family = AF_UNIX;
    strncpy(addr.sun_path, SCHCOM_SOCKET_PATH, sizeof addr.sun_path - 1);
    if (bind(lfd, (struct sockaddr *)&addr, sizeof addr) < 0)
        die("bind " SCHCOM_SOCKET_PATH);
    if (listen(lfd, 8) < 0)
        die("listen");

    int kbd = open("/dev/input/event0", O_RDONLY | O_NONBLOCK);
    int mouse = open("/dev/input/event1", O_RDONLY | O_NONBLOCK);

    printf("schcom: listening on %s\n", SCHCOM_SOCKET_PATH);
    fflush(stdout);

    composite();
    int frames = 1;

    for (;;) {
        struct pollfd pfd[MAX_CLIENTS + 3];
        int n = 0;

        pfd[n].fd = lfd;   pfd[n].events = POLLIN; n++;
        int kbd_slot = -1, mouse_slot = -1;
        if (kbd >= 0)   { kbd_slot = n;   pfd[n].fd = kbd;   pfd[n].events = POLLIN; n++; }
        if (mouse >= 0) { mouse_slot = n; pfd[n].fd = mouse; pfd[n].events = POLLIN; n++; }

        int first_client = n;
        for (int i = 0; i < MAX_CLIENTS; i++)
            if (clients[i].fd >= 0) {
                pfd[n].fd = clients[i].fd;
                pfd[n].events = POLLIN;
                n++;
            }

        if (poll(pfd, n, 100) < 0) {
            if (errno == EINTR)
                continue;
            die("poll");
        }

        int repaint = 0;

        if (pfd[0].revents & POLLIN) {
            int cfd = accept(lfd, NULL, NULL);
            if (cfd >= 0) {
                int slot = -1;
                for (int i = 0; i < MAX_CLIENTS; i++)
                    if (clients[i].fd < 0) { slot = i; break; }
                if (slot < 0) {
                    close(cfd);
                } else {
                    clients[slot].fd = cfd;
                    clients[slot].hello_done = 0;
                    printf("schcom: client %d connected\n", slot);
                    fflush(stdout);
                }
            }
        }

        if (kbd_slot >= 0 && (pfd[kbd_slot].revents & POLLIN)) {
            int r = handle_input(kbd);
            if (r < 0) break;
            repaint |= r;
        }
        if (mouse_slot >= 0 && (pfd[mouse_slot].revents & POLLIN)) {
            int r = handle_input(mouse);
            if (r < 0) break;
            repaint |= r;
        }

        /* Walk the client slots in the same order they were added to pfd. */
        int k = first_client;
        for (int i = 0; i < MAX_CLIENTS; i++) {
            if (clients[i].fd < 0)
                continue;
            if (pfd[k].revents & (POLLIN | POLLHUP))
                repaint |= handle_message(i);
            k++;
        }

        if (repaint) {
            composite();
            frames++;
            if (frames_then_exit && frames >= frames_then_exit) {
                printf("schcom: %d frames, exiting as asked\n", frames);
                fflush(stdout);
                break;
            }
        }
    }

    unlink(SCHCOM_SOCKET_PATH);
    return 0;
}