#include "kitty.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <fcntl.h>
#include <errno.h>
#include <sys/ioctl.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <signal.h>
#include <dirent.h>

// Image ids are a global namespace shared with every other program drawing in
// the same terminal, so pick something unlikely rather than 1 and 2.
#define ID_FRAME   0x51B0B01u
#define ID_OVERLAY 0x51B0B02u
// Placement ids: reusing one replaces the previous placement in-place, which is
// how we get flicker-free frame updates without deleting anything first.
#define PLACEMENT 1
// The protocol caps a single escape's payload at 4096 base64 bytes.
#define CHUNK 4096

static int zoom = 2;
static int xfer = KITTY_XFER_B64;
static bool in_tmux = false;
static bool started = false;
static int cell_w = 0, cell_h = 0;  // 0 = terminal didn't tell us

// ---- output buffer ---------------------------------------------------------
// A frame is one write(): partial graphics escapes on the wire look like line
// noise if anything else (a signal handler, a stray printf) interleaves.
static char *ob;
static size_t ob_len, ob_cap;

static void ob_need(size_t n) {
    if (ob_len + n <= ob_cap) return;
    size_t cap = ob_cap ? ob_cap : 65536;
    while (cap < ob_len + n) cap *= 2;
    char *p = realloc(ob, cap);
    if (!p) return;
    ob = p;
    ob_cap = cap;
}
static void ob_putn(const char *s, size_t n) {
    ob_need(n);
    if (ob_len + n > ob_cap) return;
    memcpy(ob + ob_len, s, n);
    ob_len += n;
}
static void ob_puts(const char *s) { ob_putn(s, strlen(s)); }

// One APC graphics command: ESC _ G <control> ; <payload> ST.
// Inside tmux the whole thing is wrapped in a passthrough DCS with every ESC
// doubled, otherwise tmux eats it (needs `set -g allow-passthrough on`).
static void apc(const char *ctrl, const char *payload, size_t plen) {
    ob_puts(in_tmux ? "\x1bPtmux;\x1b\x1b_G" : "\x1b_G");
    ob_puts(ctrl);
    if (plen) {
        ob_puts(";");
        ob_putn(payload, plen);
    }
    ob_puts(in_tmux ? "\x1b\x1b\\\x1b\\" : "\x1b\\");
}

static void ob_flush(void) {
    if (!ob_len) return;
    fwrite(ob, 1, ob_len, stdout);
    fflush(stdout);
    ob_len = 0;
}

// ---- base64 ----------------------------------------------------------------
static const char b64tab[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";

static size_t b64_len(size_t n) { return ((n + 2) / 3) * 4; }

static size_t b64_encode(const u8 *in, size_t n, char *out) {
    size_t o = 0;
    size_t i = 0;
    for (; i + 3 <= n; i += 3) {
        u32 v = ((u32)in[i] << 16) | ((u32)in[i + 1] << 8) | in[i + 2];
        out[o++] = b64tab[(v >> 18) & 63];
        out[o++] = b64tab[(v >> 12) & 63];
        out[o++] = b64tab[(v >> 6) & 63];
        out[o++] = b64tab[v & 63];
    }
    if (i < n) {
        u32 v = (u32)in[i] << 16;
        if (i + 1 < n) v |= (u32)in[i + 1] << 8;
        out[o++] = b64tab[(v >> 18) & 63];
        out[o++] = b64tab[(v >> 12) & 63];
        out[o++] = (i + 1 < n) ? b64tab[(v >> 6) & 63] : '=';
        out[o++] = '=';
    }
    return o;
}

static char *b64buf;
static size_t b64cap;
static char *b64_get(size_t need) {
    if (need > b64cap) {
        char *p = realloc(b64buf, need);
        if (!p) return NULL;
        b64buf = p;
        b64cap = need;
    }
    return b64buf;
}

// ---- terminal metrics ------------------------------------------------------
void kitty_refresh_metrics(void) {
    struct winsize ws;
    cell_w = cell_h = 0;
    if (ioctl(STDOUT_FILENO, TIOCGWINSZ, &ws) == 0 && ws.ws_col && ws.ws_row && ws.ws_xpixel &&
        ws.ws_ypixel) {
        cell_w = ws.ws_xpixel / ws.ws_col;
        cell_h = ws.ws_ypixel / ws.ws_row;
    }
}

// Does this look like a terminal that can read a shm object we created? It has
// to be the kitty/ghostty implementation *and* share our machine, so anything
// multiplexed or remote falls back to inline bytes.
static bool shm_plausible(void) {
    if (in_tmux) return false;
    const char *t = getenv("TERM");
    const char *tp = getenv("TERM_PROGRAM");
    if (getenv("KITTY_WINDOW_ID")) return true;
    if (t && strstr(t, "kitty")) return true;
    if (t && strstr(t, "ghostty")) return true;
    if (tp && (!strcmp(tp, "ghostty") || !strcmp(tp, "kitty"))) return true;
    return false;
}

static void shm_reclaim_orphans(void);  // defined with the shm transfer below

void kitty_init(int scale, int transfer) {
    if (scale < 1) scale = 1;
    if (scale > 8) scale = 8;
    zoom = scale;
    const char *tm = getenv("TMUX");
    in_tmux = tm && *tm;
    if (transfer == KITTY_XFER_SHM) xfer = KITTY_XFER_SHM;
    else if (transfer == KITTY_XFER_B64) xfer = KITTY_XFER_B64;
    else xfer = shm_plausible() ? KITTY_XFER_SHM : KITTY_XFER_B64;
    kitty_refresh_metrics();
    if (xfer == KITTY_XFER_SHM) shm_reclaim_orphans();
    started = true;
}

const char *kitty_transfer_name(void) { return xfer == KITTY_XFER_SHM ? "shm" : "b64"; }

// ---- nearest-neighbour upscale ---------------------------------------------
// Pixel art wants hard edges. The terminal scales images with linear filtering,
// so when we control the pixel count (the shm path, where bandwidth is free) we
// blow the frame up ourselves and keep it crisp.
static u8 *upbuf;
static size_t upcap;

static const u8 *upscale(const u8 *px, int w, int h, int bpp, int sc, int *ow, int *oh) {
    if (sc <= 1) {
        *ow = w;
        *oh = h;
        return px;
    }
    size_t need = (size_t)w * h * sc * sc * bpp;
    if (need > upcap) {
        u8 *p = realloc(upbuf, need);
        if (!p) {
            *ow = w;
            *oh = h;
            return px;
        }
        upbuf = p;
        upcap = need;
    }
    int W = w * sc;
    for (int y = 0; y < h; y++) {
        u8 *drow = upbuf + (size_t)y * sc * W * bpp;
        const u8 *srow = px + (size_t)y * w * bpp;
        for (int x = 0; x < w; x++) {
            const u8 *s = srow + (size_t)x * bpp;
            u8 *d = drow + (size_t)x * sc * bpp;
            for (int k = 0; k < sc; k++, d += bpp) memcpy(d, s, (size_t)bpp);
        }
        // the remaining sc-1 rows are copies of the one we just built
        for (int k = 1; k < sc; k++) memcpy(drow + (size_t)k * W * bpp, drow, (size_t)W * bpp);
    }
    *ow = W;
    *oh = h * sc;
    return upbuf;
}

// ---- pad to a whole number of cells ----------------------------------------
// The c=/r= keys size an image in *cells*, and the terminal stretches the
// pixels to fill that box - so unless the zoomed image happens to be an exact
// multiple of the cell size, the picture comes out subtly distorted (a 4:3 LCD
// going 5% tall). Growing the image by a few pixels first makes the cell box
// land exactly on zoom x, and the padding (black, or transparent for the HUD)
// is a sliver at the right/bottom edge.
static u8 *padbuf;
static size_t padcap;

static const u8 *pad_to_cells(const u8 *px, int w, int h, int bpp, int *ow, int *oh) {
    int W = w, H = h;
    while ((W * zoom) % cell_w) W++;
    while ((H * zoom) % cell_h) H++;
    *ow = W;
    *oh = H;
    if (W == w && H == h) return px;
    size_t need = (size_t)W * H * bpp;
    if (need > padcap) {
        u8 *p = realloc(padbuf, need);
        if (!p) {
            *ow = w;
            *oh = h;
            return px;
        }
        padbuf = p;
        padcap = need;
    }
    memset(padbuf, 0, need);
    for (int y = 0; y < h; y++)
        memcpy(padbuf + (size_t)y * W * bpp, px + (size_t)y * w * bpp, (size_t)w * bpp);
    return padbuf;
}

// ---- shared-memory transfer ------------------------------------------------
// The terminal unlinks the object once it has read it, so every frame needs a
// fresh name. If it *doesn't* (unsupported / it died), stale objects would pile
// up in /dev/shm, so we also unlink our own names a few frames later - harmless
// if the terminal got there first.
#define SHM_KEEP 4
static char shm_names[SHM_KEEP][64];
static unsigned shm_seq;

static bool shm_publish(const u8 *px, size_t len, char *name, size_t namesz) {
    unsigned slot = shm_seq % SHM_KEEP;
    if (shm_names[slot][0]) {  // reclaim the name from SHM_KEEP frames ago
        shm_unlink(shm_names[slot]);
        shm_names[slot][0] = 0;
    }
    snprintf(name, namesz, "/sl0pboy-%d-%u", (int)getpid(), shm_seq++);
    int fd = shm_open(name, O_RDWR | O_CREAT | O_EXCL, 0600);
    if (fd < 0) return false;
    bool ok = false;
    if (ftruncate(fd, (off_t)len) == 0) {
        void *m = mmap(NULL, len, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
        if (m != MAP_FAILED) {
            memcpy(m, px, len);
            munmap(m, len);
            ok = true;
        }
    }
    close(fd);
    if (!ok) {
        shm_unlink(name);
        return false;
    }
    snprintf(shm_names[slot], sizeof(shm_names[slot]), "%s", name);
    return true;
}

static void shm_reclaim_all(void) {
    for (int i = 0; i < SHM_KEEP; i++)
        if (shm_names[i][0]) {
            shm_unlink(shm_names[i]);
            shm_names[i][0] = 0;
        }
}

// SIGKILL leaves no chance to clean up, so an emulator that died hard can leave
// its last few objects sitting in /dev/shm forever. Sweep the ones belonging to
// pids that no longer exist before we start making our own.
static void shm_reclaim_orphans(void) {
    DIR *d = opendir("/dev/shm");
    if (!d) return;
    struct dirent *e;
    while ((e = readdir(d))) {
        int pid = 0, seq = 0;
        if (sscanf(e->d_name, "sl0pboy-%d-%d", &pid, &seq) != 2 || pid <= 0) continue;
        if (pid == getpid()) continue;
        if (kill(pid, 0) == 0 || errno != ESRCH) continue;  // still running
        char path[280];
        snprintf(path, sizeof(path), "/%s", e->d_name);
        shm_unlink(path);
    }
    closedir(d);
}

// ---- unchanged-frame suppression -------------------------------------------
// A paused game, a menu, a text box: most frames are byte-identical to the last
// one. Re-sending them is pure waste, so keep a copy per image and diff.
typedef struct {
    u32 id;
    u8 *prev;
    size_t len;
    bool valid;
} ImgCache;
static ImgCache cache[2];

static ImgCache *cache_for(u32 id) {
    for (int i = 0; i < 2; i++)
        if (cache[i].id == id) return &cache[i];
    for (int i = 0; i < 2; i++)
        if (!cache[i].id) {
            cache[i].id = id;
            return &cache[i];
        }
    return NULL;
}
// True when the pixels differ from last time (and remembers them).
static bool cache_changed(u32 id, const u8 *px, size_t len, bool force) {
    ImgCache *c = cache_for(id);
    if (!c) return true;
    if (c->len != len) {
        u8 *p = realloc(c->prev, len);
        if (!p) return true;
        c->prev = p;
        c->len = len;
        c->valid = false;
    }
    if (!force && c->valid && memcmp(c->prev, px, len) == 0) return false;
    memcpy(c->prev, px, len);
    c->valid = true;
    return true;
}
static void cache_invalidate(u32 id) {
    ImgCache *c = cache_for(id);
    if (c) c->valid = false;
}

// ---- image transmission ----------------------------------------------------
// One image, transmitted and placed in a single command (a=T). Sizing is either
// baked into the pixels (shm: we upscaled them) or delegated to the terminal
// via c=/r=, which scales the image into a cell box - that keeps the inline
// byte count at 1x no matter the zoom.
static void send_image(u32 id, const u8 *px, int w, int h, int bpp, int z, bool prescale) {
    if (!started) return;

    int sw = w, sh = h;
    const u8 *data = px;
    int cols = 0, rows = 0;

    if (xfer == KITTY_XFER_SHM || prescale) {
        data = upscale(px, w, h, bpp, zoom, &sw, &sh);
    } else if (zoom > 1) {
        if (cell_w > 0 && cell_h > 0) {
            // hand the pixels over at 1x and let the terminal scale them into a
            // cell box: the byte count stays constant whatever the zoom
            data = pad_to_cells(px, w, h, bpp, &sw, &sh);
            cols = sw * zoom / cell_w;
            rows = sh * zoom / cell_h;
        } else {
            // no cell metrics: we can't ask for a cell box, so pay the bytes
            data = upscale(px, w, h, bpp, zoom, &sw, &sh);
        }
    }

    size_t len = (size_t)sw * sh * bpp;
    char ctrl[192];
    int n = snprintf(ctrl, sizeof(ctrl), "a=T,q=2,f=%d,s=%d,v=%d,i=%u,p=%d,C=1", bpp == 4 ? 32 : 24,
                     sw, sh, id, PLACEMENT);
    if (z) n += snprintf(ctrl + n, sizeof(ctrl) - n, ",z=%d", z);
    if (cols) n += snprintf(ctrl + n, sizeof(ctrl) - n, ",c=%d,r=%d", cols, rows);

    if (xfer == KITTY_XFER_SHM) {
        char name[64];
        if (shm_publish(data, len, name, sizeof(name))) {
            char nb64[128];
            size_t nn = b64_encode((const u8 *)name, strlen(name), nb64);
            snprintf(ctrl + n, sizeof(ctrl) - n, ",t=s");
            apc(ctrl, nb64, nn);
            return;
        }
        // shm went away under us (no /dev/shm?) - inline the bytes instead
        xfer = KITTY_XFER_B64;
        cache_invalidate(id);
    }

    size_t need = b64_len(len);
    char *enc = b64_get(need);
    if (!enc) return;
    size_t elen = b64_encode(data, len, enc);

    // The payload rides in <=4096-byte base64 chunks: the first carries the
    // control data and m=1 ("more coming"), the rest only need m, and m=0 ends
    // the transmission.
    size_t off = 0;
    bool first = true;
    while (first || off < elen) {
        size_t take = elen - off < CHUNK ? elen - off : CHUNK;
        bool last = (off + take >= elen);
        if (first) {
            char full[224];
            snprintf(full, sizeof(full), "%s,t=d,m=%d", ctrl, last ? 0 : 1);
            apc(full, enc + off, take);
            first = false;
        } else {
            apc(last ? "m=0" : "m=1", enc + off, take);
        }
        off += take;
    }
}

void kitty_frame_rgb(const u8 *rgb, int w, int h, bool force) {
    size_t len = (size_t)w * h * 3;
    if (!cache_changed(ID_FRAME, rgb, len, force)) return;
    ob_puts("\x1b[H");  // place at the top-left cell (never wrapped:
                        // tmux has to see this one to track the pane)
    send_image(ID_FRAME, rgb, w, h, 3, 0, false);
    ob_flush();
}

void kitty_overlay_rgba(const u8 *rgba, int w, int h, bool force) {
    size_t len = (size_t)w * h * 4;
    if (!cache_changed(ID_OVERLAY, rgba, len, force)) return;
    ob_puts("\x1b[H");
    // The HUD is small and only changes when its text does, so it always ships
    // pre-zoomed: no cell-box rounding to distort a 22px-tall panel.
    send_image(ID_OVERLAY, rgba, w, h, 4, 1, true);  // z=1: above the frame
    ob_flush();
}

void kitty_overlay_clear(void) {
    if (!started) return;
    char ctrl[64];
    snprintf(ctrl, sizeof(ctrl), "a=d,d=I,i=%u,q=2", ID_OVERLAY);
    apc(ctrl, NULL, 0);
    ob_flush();
    cache_invalidate(ID_OVERLAY);
}

void kitty_shutdown(void) {
    if (!started) return;
    char ctrl[64];
    snprintf(ctrl, sizeof(ctrl), "a=d,d=I,i=%u,q=2", ID_FRAME);
    apc(ctrl, NULL, 0);
    snprintf(ctrl, sizeof(ctrl), "a=d,d=I,i=%u,q=2", ID_OVERLAY);
    apc(ctrl, NULL, 0);
    ob_flush();
    shm_reclaim_all();
    started = false;
}