#include "session.h"

#include <dirent.h>
#include <errno.h>
#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include <sys/stat.h>
#include <sys/types.h>

#include <termios.h>

#define LOG_MODULE "session"
#define LOG_ENABLE_DBG 0
#include "log.h"
#include "xmalloc.h"

#define JSON_SUFFIX ".json"
#define JSON_SUFFIX_LEN 5
#define SCROLLBACK_SUFFIX ".scrollback.txt"
#define SCROLLBACK_ENC_SUFFIX ".scrollback.enc"
#define SCROLLBACK_MAX_BYTES (1u << 20)  /* 1 MB cap */

static int
state_dir(char *buf, size_t buf_sz)
{
    const char *xdg = getenv("XDG_DATA_HOME");
    int n;
    if (xdg != NULL && xdg[0] == '/')
        n = snprintf(buf, buf_sz, "%s/foot/state", xdg);
    else {
        const char *home = getenv("HOME");
        if (home == NULL)
            return -1;
        n = snprintf(buf, buf_sz, "%s/.local/share/foot/state", home);
    }
    if (n < 0 || (size_t)n >= buf_sz)
        return -1;
    return n;
}

static bool
mkdir_p(const char *path)
{
    char tmp[4096];
    size_t len = strlen(path);
    if (len >= sizeof(tmp))
        return false;
    memcpy(tmp, path, len + 1);

    for (size_t i = 1; i < len; i++) {
        if (tmp[i] == '/') {
            tmp[i] = '\0';
            if (mkdir(tmp, 0700) < 0 && errno != EEXIST) {
                LOG_ERRNO("mkdir(%s)", tmp);
                return false;
            }
            tmp[i] = '/';
        }
    }
    if (mkdir(tmp, 0700) < 0 && errno != EEXIST) {
        LOG_ERRNO("mkdir(%s)", tmp);
        return false;
    }
    return true;
}

static int
session_path(const char *name, char *buf, size_t buf_sz)
{
    char dir[1024];
    if (state_dir(dir, sizeof(dir)) < 0)
        return -1;
    int n = snprintf(buf, buf_sz, "%s/%s.json", dir, name);
    if (n < 0 || (size_t)n >= buf_sz)
        return -1;
    return n;
}

bool
session_name_is_valid(const char *name)
{
    if (name == NULL || name[0] == '\0')
        return false;
    for (const char *p = name; *p != '\0'; p++) {
        char c = *p;
        bool ok = (c >= 'A' && c <= 'Z') || (c >= 'a' && c <= 'z') ||
                  (c >= '0' && c <= '9') || c == '.' || c == '_' || c == '-';
        if (!ok)
            return false;
    }
    return true;
}

void
session_state_free(struct session_state *st)
{
    if (st == NULL)
        return;
    free(st->cwd);
    if (st->argv != NULL) {
        for (int i = 0; i < st->argc; i++)
            free(st->argv[i]);
        free(st->argv);
    }
    st->cwd = NULL;
    st->argv = NULL;
    st->argc = 0;
}

#if defined(__linux__)
static char *
read_link(const char *path)
{
    char *buf = NULL;
    size_t cap = 256;
    for (;;) {
        buf = xrealloc(buf, cap);
        ssize_t n = readlink(path, buf, cap - 1);
        if (n < 0) {
            free(buf);
            return NULL;
        }
        if ((size_t)n < cap - 1) {
            buf[n] = '\0';
            return buf;
        }
        cap *= 2;
        if (cap > 1 << 16) {
            free(buf);
            return NULL;
        }
    }
}

static bool
read_cmdline(pid_t pid, int *out_argc, char ***out_argv)
{
    char path[64];
    snprintf(path, sizeof(path), "/proc/%d/cmdline", (int)pid);

    int fd = open(path, O_RDONLY | O_CLOEXEC);
    if (fd < 0)
        return false;

    char *buf = NULL;
    size_t len = 0;
    size_t cap = 0;
    char chunk[1024];
    for (;;) {
        ssize_t n = read(fd, chunk, sizeof(chunk));
        if (n < 0) {
            if (errno == EINTR)
                continue;
            close(fd);
            free(buf);
            return false;
        }
        if (n == 0)
            break;
        if (len + (size_t)n + 1 > cap) {
            cap = (cap == 0 ? 1024 : cap * 2);
            while (len + (size_t)n + 1 > cap)
                cap *= 2;
            buf = xrealloc(buf, cap);
        }
        memcpy(buf + len, chunk, n);
        len += n;
    }
    close(fd);

    if (buf == NULL || len == 0) {
        free(buf);
        return false;
    }
    /* /proc/<pid>/cmdline is NUL-separated; ensure trailing NUL */
    if (buf[len - 1] != '\0') {
        buf = xrealloc(buf, len + 1);
        buf[len] = '\0';
        len++;
    }

    /* Count args */
    int argc = 0;
    for (size_t i = 0; i < len; i++)
        if (buf[i] == '\0' && (i == 0 || buf[i - 1] != '\0'))
            argc++;

    if (argc <= 0) {
        free(buf);
        return false;
    }

    char **argv = xmalloc(sizeof(char *) * (argc + 1));
    int idx = 0;
    size_t start = 0;
    for (size_t i = 0; i < len && idx < argc; i++) {
        if (buf[i] == '\0') {
            argv[idx++] = xstrdup(buf + start);
            start = i + 1;
        }
    }
    argv[argc] = NULL;
    free(buf);

    *out_argc = argc;
    *out_argv = argv;
    return true;
}
#endif

bool
session_capture(int ptmx, const char *fallback_cwd, struct session_state *out)
{
    memset(out, 0, sizeof(*out));

#if defined(__linux__)
    pid_t fg = tcgetpgrp(ptmx);
    if (fg > 0) {
        char path[64];
        snprintf(path, sizeof(path), "/proc/%d/cwd", (int)fg);
        char *cwd = read_link(path);
        if (cwd != NULL)
            out->cwd = cwd;

        int argc = 0;
        char **argv = NULL;
        if (read_cmdline(fg, &argc, &argv)) {
            /* If the foreground process is a shell (login or otherwise), don't
             * try to resurrect it via argv — let the user's normal shell
             * launch take over. We only restore non-shell processes. */
            const char *base = argv[0];
            const char *slash = strrchr(base, '/');
            if (slash != NULL)
                base = slash + 1;
            const char *shells[] = {
                "bash", "zsh", "fish", "sh", "dash", "ksh", "tcsh", "csh",
                "-bash", "-zsh", "-fish", "-sh", "-dash"
            };
            bool is_shell = false;
            for (size_t i = 0; i < sizeof(shells) / sizeof(shells[0]); i++) {
                if (strcmp(base, shells[i]) == 0) {
                    is_shell = true;
                    break;
                }
            }
            if (is_shell) {
                for (int i = 0; i < argc; i++)
                    free(argv[i]);
                free(argv);
            } else {
                out->argc = argc;
                out->argv = argv;
            }
        }
    }
#endif

    if (out->cwd == NULL && fallback_cwd != NULL)
        out->cwd = xstrdup(fallback_cwd);

    return out->cwd != NULL;
}

/* ------------- JSON encoding/decoding (minimal, hand-rolled) ------------- */

static void
json_escape(FILE *fp, const char *s)
{
    fputc('"', fp);
    for (const char *p = s; *p != '\0'; p++) {
        unsigned char c = (unsigned char)*p;
        switch (c) {
        case '"': fputs("\\\"", fp); break;
        case '\\': fputs("\\\\", fp); break;
        case '\b': fputs("\\b", fp); break;
        case '\f': fputs("\\f", fp); break;
        case '\n': fputs("\\n", fp); break;
        case '\r': fputs("\\r", fp); break;
        case '\t': fputs("\\t", fp); break;
        default:
            if (c < 0x20)
                fprintf(fp, "\\u%04x", c);
            else
                fputc(c, fp);
        }
    }
    fputc('"', fp);
}

bool
session_save(const char *name, const struct session_state *st)
{
    if (!session_name_is_valid(name)) {
        LOG_ERR("invalid session name: %s", name == NULL ? "(null)" : name);
        return false;
    }
    if (st == NULL || st->cwd == NULL) {
        LOG_ERR("session_save: missing cwd");
        return false;
    }

    char dir[1024];
    if (state_dir(dir, sizeof(dir)) < 0) {
        LOG_ERR("could not resolve state directory");
        return false;
    }
    if (!mkdir_p(dir))
        return false;

    char path[1280];
    if (session_path(name, path, sizeof(path)) < 0)
        return false;

    char tmp_path[1408];
    snprintf(tmp_path, sizeof(tmp_path), "%s.tmp", path);

    FILE *fp = fopen(tmp_path, "we");
    if (fp == NULL) {
        LOG_ERRNO("fopen(%s)", tmp_path);
        return false;
    }

    fputs("{\n  \"cwd\": ", fp);
    json_escape(fp, st->cwd);
    if (st->argc > 0 && st->argv != NULL) {
        fputs(",\n  \"argv\": [", fp);
        for (int i = 0; i < st->argc; i++) {
            if (i > 0)
                fputs(", ", fp);
            json_escape(fp, st->argv[i]);
        }
        fputc(']', fp);
    }
    fputs("\n}\n", fp);

    if (fflush(fp) != 0 || fsync(fileno(fp)) != 0) {
        LOG_ERRNO("fsync(%s)", tmp_path);
        fclose(fp);
        unlink(tmp_path);
        return false;
    }
    fclose(fp);

    if (rename(tmp_path, path) < 0) {
        LOG_ERRNO("rename(%s -> %s)", tmp_path, path);
        unlink(tmp_path);
        return false;
    }

    LOG_INFO("session saved: %s", path);
    return true;
}

/* ------------- JSON parsing (just enough for our schema) ------------- */

struct parser {
    const char *p;
    const char *end;
};

static void
skip_ws(struct parser *pa)
{
    while (pa->p < pa->end) {
        char c = *pa->p;
        if (c == ' ' || c == '\t' || c == '\n' || c == '\r')
            pa->p++;
        else
            break;
    }
}

static bool
parse_string(struct parser *pa, char **out)
{
    skip_ws(pa);
    if (pa->p >= pa->end || *pa->p != '"')
        return false;
    pa->p++;

    size_t cap = 64;
    size_t len = 0;
    char *buf = xmalloc(cap);

    while (pa->p < pa->end && *pa->p != '"') {
        char c = *pa->p++;
        char decoded;
        if (c == '\\') {
            if (pa->p >= pa->end) { free(buf); return false; }
            char esc = *pa->p++;
            switch (esc) {
            case '"': decoded = '"'; break;
            case '\\': decoded = '\\'; break;
            case '/': decoded = '/'; break;
            case 'b': decoded = '\b'; break;
            case 'f': decoded = '\f'; break;
            case 'n': decoded = '\n'; break;
            case 'r': decoded = '\r'; break;
            case 't': decoded = '\t'; break;
            case 'u': {
                if (pa->end - pa->p < 4) { free(buf); return false; }
                unsigned val = 0;
                for (int i = 0; i < 4; i++) {
                    char h = pa->p[i];
                    val <<= 4;
                    if (h >= '0' && h <= '9') val |= h - '0';
                    else if (h >= 'a' && h <= 'f') val |= h - 'a' + 10;
                    else if (h >= 'A' && h <= 'F') val |= h - 'A' + 10;
                    else { free(buf); return false; }
                }
                pa->p += 4;
                if (val < 0x80) {
                    decoded = (char)val;
                } else {
                    /* UTF-8 encode codepoint (no surrogate-pair handling) */
                    char enc[4];
                    int n;
                    if (val < 0x800) {
                        enc[0] = (char)(0xc0 | (val >> 6));
                        enc[1] = (char)(0x80 | (val & 0x3f));
                        n = 2;
                    } else {
                        enc[0] = (char)(0xe0 | (val >> 12));
                        enc[1] = (char)(0x80 | ((val >> 6) & 0x3f));
                        enc[2] = (char)(0x80 | (val & 0x3f));
                        n = 3;
                    }
                    if (len + (size_t)n + 1 > cap) {
                        while (len + (size_t)n + 1 > cap) cap *= 2;
                        buf = xrealloc(buf, cap);
                    }
                    memcpy(buf + len, enc, n);
                    len += n;
                    continue;
                }
                break;
            }
            default: free(buf); return false;
            }
        } else {
            decoded = c;
        }
        if (len + 2 > cap) {
            cap *= 2;
            buf = xrealloc(buf, cap);
        }
        buf[len++] = decoded;
    }
    if (pa->p >= pa->end) { free(buf); return false; }
    pa->p++;  /* consume closing quote */
    buf[len] = '\0';
    *out = buf;
    return true;
}

static bool
expect_char(struct parser *pa, char c)
{
    skip_ws(pa);
    if (pa->p >= pa->end || *pa->p != c)
        return false;
    pa->p++;
    return true;
}

bool
session_load(const char *name, struct session_state *out)
{
    memset(out, 0, sizeof(*out));

    if (!session_name_is_valid(name)) {
        LOG_ERR("invalid session name: %s", name == NULL ? "(null)" : name);
        return false;
    }

    char path[1280];
    if (session_path(name, path, sizeof(path)) < 0)
        return false;

    int fd = open(path, O_RDONLY | O_CLOEXEC);
    if (fd < 0) {
        LOG_ERRNO("open(%s)", path);
        return false;
    }
    struct stat sb;
    if (fstat(fd, &sb) < 0 || sb.st_size <= 0 || sb.st_size > (off_t)(1 << 20)) {
        LOG_ERR("session file %s has unexpected size", path);
        close(fd);
        return false;
    }
    char *data = xmalloc(sb.st_size + 1);
    ssize_t got = 0;
    while (got < sb.st_size) {
        ssize_t n = read(fd, data + got, sb.st_size - got);
        if (n < 0) {
            if (errno == EINTR) continue;
            free(data);
            close(fd);
            return false;
        }
        if (n == 0) break;
        got += n;
    }
    close(fd);
    data[got] = '\0';

    struct parser pa = { .p = data, .end = data + got };
    bool ok = false;

    if (!expect_char(&pa, '{')) goto done;

    for (;;) {
        skip_ws(&pa);
        if (pa.p < pa.end && *pa.p == '}') { pa.p++; break; }

        char *key = NULL;
        if (!parse_string(&pa, &key)) goto done;
        if (!expect_char(&pa, ':')) { free(key); goto done; }

        if (strcmp(key, "cwd") == 0) {
            free(key);
            if (!parse_string(&pa, &out->cwd)) goto done;
        } else if (strcmp(key, "argv") == 0) {
            free(key);
            if (!expect_char(&pa, '[')) goto done;
            int cap = 4;
            out->argv = xmalloc(sizeof(char *) * cap);
            out->argc = 0;
            for (;;) {
                skip_ws(&pa);
                if (pa.p < pa.end && *pa.p == ']') { pa.p++; break; }
                if (out->argc + 1 >= cap) {
                    cap *= 2;
                    out->argv = xrealloc(out->argv, sizeof(char *) * cap);
                }
                if (!parse_string(&pa, &out->argv[out->argc])) goto done;
                out->argc++;
                skip_ws(&pa);
                if (pa.p < pa.end && *pa.p == ',') { pa.p++; continue; }
            }
            out->argv[out->argc] = NULL;
        } else {
            free(key);
            /* unknown key: skip a string value */
            char *junk = NULL;
            if (!parse_string(&pa, &junk)) goto done;
            free(junk);
        }

        skip_ws(&pa);
        if (pa.p < pa.end && *pa.p == ',') { pa.p++; continue; }
    }

    ok = out->cwd != NULL;

done:
    free(data);
    if (!ok)
        session_state_free(out);
    return ok;
}

/* ------------- listing ------------- */

static int
cmp_names(const void *a, const void *b)
{
    return strcmp(*(const char *const *)a, *(const char *const *)b);
}

char **
session_list(size_t *out_count)
{
    *out_count = 0;

    char dir[1024];
    if (state_dir(dir, sizeof(dir)) < 0)
        return NULL;

    DIR *d = opendir(dir);
    if (d == NULL)
        return NULL;

    size_t cap = 8;
    size_t count = 0;
    char **names = xmalloc(sizeof(char *) * cap);

    struct dirent *de;
    while ((de = readdir(d)) != NULL) {
        const char *n = de->d_name;
        size_t nlen = strlen(n);
        if (nlen <= JSON_SUFFIX_LEN)
            continue;
        if (strcmp(n + nlen - JSON_SUFFIX_LEN, JSON_SUFFIX) != 0)
            continue;
        char *stripped = xmalloc(nlen - JSON_SUFFIX_LEN + 1);
        memcpy(stripped, n, nlen - JSON_SUFFIX_LEN);
        stripped[nlen - JSON_SUFFIX_LEN] = '\0';
        if (!session_name_is_valid(stripped)) {
            free(stripped);
            continue;
        }
        if (count >= cap) {
            cap *= 2;
            names = xrealloc(names, sizeof(char *) * cap);
        }
        names[count++] = stripped;
    }
    closedir(d);

    qsort(names, count, sizeof(char *), cmp_names);

    *out_count = count;
    return names;
}

bool
session_exists(const char *name)
{
    if (!session_name_is_valid(name))
        return false;
    char path[1280];
    if (session_path(name, path, sizeof(path)) < 0)
        return false;
    return access(path, F_OK) == 0;
}

bool
session_delete(const char *name)
{
    if (!session_name_is_valid(name))
        return false;
    char path[1280];
    if (session_path(name, path, sizeof(path)) < 0)
        return false;
    if (unlink(path) < 0) {
        LOG_ERRNO("unlink(%s)", path);
        return false;
    }
    session_delete_scrollback(name);
    session_delete_enc_scrollback(name);
    LOG_INFO("session deleted: %s", path);
    return true;
}

static int
scrollback_path(const char *name, char *buf, size_t buf_sz)
{
    char dir[1024];
    if (state_dir(dir, sizeof(dir)) < 0)
        return -1;
    int n = snprintf(buf, buf_sz, "%s/%s%s", dir, name, SCROLLBACK_SUFFIX);
    if (n < 0 || (size_t)n >= buf_sz)
        return -1;
    return n;
}

bool
session_load_scrollback(const char *name, char **out_text, size_t *out_len)
{
    *out_text = NULL;
    *out_len = 0;
    if (!session_name_is_valid(name))
        return false;

    char path[1280];
    if (scrollback_path(name, path, sizeof(path)) < 0)
        return false;

    int fd = open(path, O_RDONLY | O_CLOEXEC);
    if (fd < 0)
        return false;

    struct stat sb;
    if (fstat(fd, &sb) < 0 || sb.st_size < 0 ||
        (size_t)sb.st_size > SCROLLBACK_MAX_BYTES * 4 /* generous on load */) {
        close(fd);
        return false;
    }
    if (sb.st_size == 0) {
        close(fd);
        return false;
    }

    char *data = xmalloc(sb.st_size + 1);
    ssize_t got = 0;
    while (got < sb.st_size) {
        ssize_t n = read(fd, data + got, sb.st_size - got);
        if (n < 0) {
            if (errno == EINTR) continue;
            free(data);
            close(fd);
            return false;
        }
        if (n == 0) break;
        got += n;
    }
    close(fd);
    data[got] = '\0';

    *out_text = data;
    *out_len = (size_t)got;
    return true;
}

void
session_delete_scrollback(const char *name)
{
    if (!session_name_is_valid(name))
        return;
    char path[1280];
    if (scrollback_path(name, path, sizeof(path)) < 0)
        return;
    unlink(path);
}

static int
enc_path(const char *name, char *buf, size_t buf_sz)
{
    char dir[1024];
    if (state_dir(dir, sizeof(dir)) < 0)
        return -1;
    int n = snprintf(buf, buf_sz, "%s/%s%s", dir, name, SCROLLBACK_ENC_SUFFIX);
    if (n < 0 || (size_t)n >= buf_sz)
        return -1;
    return n;
}

bool
session_has_enc_scrollback(const char *name)
{
    if (!session_name_is_valid(name))
        return false;
    char path[1280];
    if (enc_path(name, path, sizeof(path)) < 0)
        return false;
    return access(path, F_OK) == 0;
}

void
session_delete_enc_scrollback(const char *name)
{
    if (!session_name_is_valid(name))
        return;
    char path[1280];
    if (enc_path(name, path, sizeof(path)) < 0)
        return;
    unlink(path);
}

bool
session_write_enc_blob(const char *name, const unsigned char *blob, size_t len)
{
    if (!session_name_is_valid(name))
        return false;
    char dir[1024];
    if (state_dir(dir, sizeof(dir)) < 0)
        return false;
    if (!mkdir_p(dir))
        return false;
    char path[1280];
    if (enc_path(name, path, sizeof(path)) < 0)
        return false;
    char tmp_path[1408];
    snprintf(tmp_path, sizeof(tmp_path), "%s.tmp", path);

    /* 0600 — encrypted but still reduces accidental exposure. */
    int fd = open(tmp_path, O_WRONLY | O_CREAT | O_TRUNC | O_CLOEXEC, 0600);
    if (fd < 0) {
        LOG_ERRNO("open(%s)", tmp_path);
        return false;
    }
    size_t written = 0;
    while (written < len) {
        ssize_t n = write(fd, blob + written, len - written);
        if (n < 0) {
            if (errno == EINTR) continue;
            LOG_ERRNO("write(%s)", tmp_path);
            close(fd);
            unlink(tmp_path);
            return false;
        }
        written += n;
    }
    if (fsync(fd) < 0) {
        LOG_ERRNO("fsync(%s)", tmp_path);
        close(fd);
        unlink(tmp_path);
        return false;
    }
    close(fd);
    if (rename(tmp_path, path) < 0) {
        LOG_ERRNO("rename(%s -> %s)", tmp_path, path);
        unlink(tmp_path);
        return false;
    }
    return true;
}

bool
session_read_enc_blob(const char *name, unsigned char **out, size_t *out_len)
{
    *out = NULL;
    *out_len = 0;
    if (!session_name_is_valid(name))
        return false;
    char path[1280];
    if (enc_path(name, path, sizeof(path)) < 0)
        return false;
    int fd = open(path, O_RDONLY | O_CLOEXEC);
    if (fd < 0)
        return false;
    struct stat sb;
    if (fstat(fd, &sb) < 0 || sb.st_size <= 0 ||
        (size_t)sb.st_size > SCROLLBACK_MAX_BYTES * 4) {
        close(fd);
        return false;
    }
    unsigned char *buf = xmalloc(sb.st_size);
    ssize_t got = 0;
    while (got < sb.st_size) {
        ssize_t n = read(fd, buf + got, sb.st_size - got);
        if (n < 0) {
            if (errno == EINTR) continue;
            free(buf);
            close(fd);
            return false;
        }
        if (n == 0) break;
        got += n;
    }
    close(fd);
    if (got != sb.st_size) {
        free(buf);
        return false;
    }
    *out = buf;
    *out_len = (size_t)got;
    return true;
}

void
session_free_names(char **names, size_t count)
{
    if (names == NULL)
        return;
    for (size_t i = 0; i < count; i++)
        free(names[i]);
    free(names);
}