/*
 * orionlock — PAM-authenticated fullscreen screen lock.
 * Grabs keyboard + pointer, draws clock and password entry.
 * On correct password, ungrabs and exits.
 */
#include <X11/Xlib.h>
#include <X11/Xutil.h>
#include <X11/keysym.h>
#include <X11/Xft/Xft.h>
#include <security/pam_appl.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <time.h>
#include <unistd.h>

#define COL_BG      0x000000
#define COL_CLOCK   0xFFFFFF
#define COL_DATE    0x8899AA
#define COL_DOT     0x4A90D9
#define COL_ERR     0xFF5F57
#define COL_HINT    0x555566
#define PW_MAX      256
#define DOT_R       5
#define DOT_PAD     8

static Display  *dpy;
static int       screen;
static Window    root, lock_win;
static Visual   *visual;
static Colormap  cmap;
static GC        gc;
static XftFont  *font_clock;
static XftFont  *font_date;
static XftFont  *font_ui;

static char      pw_buf[PW_MAX];
static int       pw_len    = 0;
static int       auth_fail = 0;
static int       auth_ok   = 0;

static unsigned long alloc_color(unsigned long rgb) {
    XColor c;
    c.red   = ((rgb >> 16) & 0xFF) * 257;
    c.green = ((rgb >>  8) & 0xFF) * 257;
    c.blue  = ((rgb)       & 0xFF) * 257;
    c.flags = DoRed | DoGreen | DoBlue;
    XAllocColor(dpy, cmap, &c);
    return c.pixel;
}

static void xft_color(unsigned long rgb, XftColor *out) {
    XRenderColor rc;
    rc.red   = ((rgb >> 16) & 0xFF) * 257;
    rc.green = ((rgb >>  8) & 0xFF) * 257;
    rc.blue  = ((rgb)       & 0xFF) * 257;
    rc.alpha = 0xFFFF;
    XftColorAllocValue(dpy, visual, cmap, &rc, out);
}

static void draw_lock(void) {
    int sw = DisplayWidth(dpy, screen);
    int sh = DisplayHeight(dpy, screen);

    XSetForeground(dpy, gc, alloc_color(COL_BG));
    XFillRectangle(dpy, lock_win, gc, 0, 0, sw, sh);

    XftDraw *xd = XftDrawCreate(dpy, lock_win, visual, cmap);
    XftColor col;
    XGlyphInfo ext;

    /* clock */
    {
        time_t t = time(NULL);
        struct tm *tm = localtime(&t);
        char clock_str[16], date_str[32];
        strftime(clock_str, sizeof(clock_str), "%H:%M", tm);
        strftime(date_str,  sizeof(date_str),  "%A, %B %d", tm);

        xft_color(COL_CLOCK, &col);
        XftTextExtentsUtf8(dpy, font_clock, (const FcChar8 *)clock_str,
                           strlen(clock_str), &ext);
        int cx = (sw - (int)ext.width) / 2;
        int cy = sh / 2 - 100;
        XftDrawStringUtf8(xd, &col, font_clock, cx, cy,
                          (const FcChar8 *)clock_str, strlen(clock_str));
        XftColorFree(dpy, visual, cmap, &col);

        xft_color(COL_DATE, &col);
        XftTextExtentsUtf8(dpy, font_date, (const FcChar8 *)date_str,
                           strlen(date_str), &ext);
        cx = (sw - (int)ext.width) / 2;
        XftDrawStringUtf8(xd, &col, font_date, cx, cy + 40,
                          (const FcChar8 *)date_str, strlen(date_str));
        XftColorFree(dpy, visual, cmap, &col);
    }

    /* password dots */
    {
        int dots = pw_len;
        if (dots > 20) dots = 20;
        int total_w = dots * (DOT_R * 2 + DOT_PAD) - (dots > 0 ? DOT_PAD : 0);
        int start_x = (sw - total_w) / 2;
        int dot_y   = sh / 2 + 10;
        unsigned long dot_col = alloc_color(auth_fail ? COL_ERR : COL_DOT);
        XSetForeground(dpy, gc, dot_col);
        for (int i = 0; i < dots; i++) {
            int dx = start_x + i * (DOT_R * 2 + DOT_PAD);
            XFillArc(dpy, lock_win, gc,
                     dx, dot_y - DOT_R, DOT_R * 2, DOT_R * 2, 0, 360 * 64);
        }
        if (dots == 0 && pw_len == 0) {
            /* hint: show lock icon placeholder */
            xft_color(COL_HINT, &col);
            const char *hint = "Type password to unlock";
            XftTextExtentsUtf8(dpy, font_ui, (const FcChar8 *)hint,
                               strlen(hint), &ext);
            int hx = (sw - (int)ext.width) / 2;
            XftDrawStringUtf8(xd, &col, font_ui, hx, dot_y + DOT_R + 18,
                              (const FcChar8 *)hint, strlen(hint));
            XftColorFree(dpy, visual, cmap, &col);
        }
    }

    /* error message */
    if (auth_fail) {
        xft_color(COL_ERR, &col);
        const char *msg = "Incorrect password";
        XftTextExtentsUtf8(dpy, font_ui, (const FcChar8 *)msg,
                           strlen(msg), &ext);
        int ex = (sw - (int)ext.width) / 2;
        XftDrawStringUtf8(xd, &col, font_ui, ex, sh / 2 + 60,
                          (const FcChar8 *)msg, strlen(msg));
        XftColorFree(dpy, visual, cmap, &col);
    }

    XftDrawDestroy(xd);
    XFlush(dpy);
}

/* ── PAM ─────────────────────────────────────────────────────────────── */

static const char *pam_pw = NULL;

static int pam_conv_fn(int num_msg, const struct pam_message **msg,
                        struct pam_response **resp, void *appdata_ptr) {
    (void)appdata_ptr;
    struct pam_response *r = calloc(num_msg, sizeof(struct pam_response));
    if (!r) return PAM_CONV_ERR;
    for (int i = 0; i < num_msg; i++) {
        if (msg[i]->msg_style == PAM_PROMPT_ECHO_OFF ||
            msg[i]->msg_style == PAM_PROMPT_ECHO_ON) {
            r[i].resp = strdup(pam_pw ? pam_pw : "");
        }
    }
    *resp = r;
    return PAM_SUCCESS;
}

static int try_auth(const char *password) {
    const char *user = getenv("USER");
    if (!user) user = "root";
    pam_pw = password;

    struct pam_conv conv = { pam_conv_fn, NULL };
    pam_handle_t *ph = NULL;
    int ret = pam_start("login", user, &conv, &ph);
    if (ret != PAM_SUCCESS) { pam_end(ph, ret); return 0; }
    ret = pam_authenticate(ph, PAM_SILENT);
    pam_end(ph, ret);
    return (ret == PAM_SUCCESS);
}

int main(void) {
    dpy    = XOpenDisplay(NULL);
    if (!dpy) { fprintf(stderr, "orionlock: cannot open display\n"); return 1; }
    screen = DefaultScreen(dpy);
    root   = RootWindow(dpy, screen);
    visual = DefaultVisual(dpy, screen);
    cmap   = DefaultColormap(dpy, screen);

    int sw = DisplayWidth(dpy, screen);
    int sh = DisplayHeight(dpy, screen);

    XSetWindowAttributes swa;
    swa.override_redirect = True;
    swa.background_pixel  = alloc_color(COL_BG);
    swa.event_mask        = ExposureMask | KeyPressMask;
    lock_win = XCreateWindow(dpy, root, 0, 0, sw, sh, 0,
                              DefaultDepth(dpy, screen), InputOutput,
                              visual, CWOverrideRedirect | CWBackPixel | CWEventMask, &swa);
    gc = XCreateGC(dpy, lock_win, 0, NULL);

    font_clock = XftFontOpenName(dpy, screen, "Sans:size=64:weight=bold");
    font_date  = XftFontOpenName(dpy, screen, "Sans:size=20");
    font_ui    = XftFontOpenName(dpy, screen, "Sans:size=11");
    if (!font_clock) font_clock = XftFontOpenName(dpy, screen, "fixed:size=20");
    if (!font_date)  font_date  = font_clock;
    if (!font_ui)    font_ui    = font_clock;

    XMapRaised(dpy, lock_win);
    XSync(dpy, False);

    /* grab keyboard and pointer */
    for (int t = 0; t < 10; t++) {
        int r = XGrabKeyboard(dpy, lock_win, True, GrabModeAsync, GrabModeAsync, CurrentTime);
        if (r == GrabSuccess) break;
        usleep(100000);
    }
    XGrabPointer(dpy, lock_win, True, ButtonPressMask | ButtonReleaseMask |
                 PointerMotionMask, GrabModeAsync, GrabModeAsync,
                 lock_win, None, CurrentTime);

    draw_lock();

    time_t last_tick = 0;

    while (!auth_ok) {
        /* update clock every second even with no events */
        fd_set fds;
        FD_ZERO(&fds);
        FD_SET(ConnectionNumber(dpy), &fds);
        struct timeval tv = { 1, 0 };
        select(ConnectionNumber(dpy) + 1, &fds, NULL, NULL, &tv);

        time_t now = time(NULL);
        if (now != last_tick) { last_tick = now; draw_lock(); }

        while (XPending(dpy)) {
            XEvent e;
            XNextEvent(dpy, &e);
            if (e.type == Expose && e.xexpose.count == 0) {
                draw_lock();
            } else if (e.type == KeyPress) {
                char buf[32] = {0};
                KeySym sym;
                XLookupString(&e.xkey, buf, sizeof(buf) - 1, &sym, NULL);

                if (sym == XK_Return || sym == XK_KP_Enter) {
                    pw_buf[pw_len] = '\0';
                    if (try_auth(pw_buf)) {
                        auth_ok = 1;
                    } else {
                        auth_fail = 1;
                        pw_len = 0;
                        draw_lock();
                        usleep(400000);
                        auth_fail = 0;
                    }
                } else if (sym == XK_BackSpace || sym == XK_Delete) {
                    if (pw_len > 0) pw_len--;
                    auth_fail = 0;
                } else if (sym == XK_Escape) {
                    pw_len    = 0;
                    auth_fail = 0;
                } else if (buf[0] >= 0x20 && pw_len < PW_MAX - 1) {
                    pw_buf[pw_len++] = buf[0];
                    auth_fail = 0;
                }
                draw_lock();
            }
        }
    }

    XUngrabKeyboard(dpy, CurrentTime);
    XUngrabPointer(dpy, CurrentTime);
    XDestroyWindow(dpy, lock_win);
    XSync(dpy, False);
    XCloseDisplay(dpy);
    return 0;
}