#include "limits.h"

#include <ctype.h>
#include <stdlib.h>
#include <string.h>

#define DEFAULT_SHELL_TIMEOUT_SECS 120UL
#define DEFAULT_MAX_OUTPUT_CHARS 16000UL
#define DEFAULT_MAX_FILE_BYTES 2000000UL

static tool_limits g_limits;
static bool g_limits_inited = false;
static bool (*g_confirm_hook)(const char *, void *) = NULL;
static void *g_confirm_ud = NULL;

/* The global starts zeroed in BSS; make sure it holds defaults before
 * any consumer reads it (Rust: lazy static initialized to defaults). */
static void limits_ensure(void) {
    if (!g_limits_inited) {
        tool_limits_defaults(&g_limits);
        g_limits_inited = true;
    }
}

void tool_limits_defaults(tool_limits *l) {
    memset(l, 0, sizeof *l);
    l->shell_timeout_secs = DEFAULT_SHELL_TIMEOUT_SECS;
    l->max_output_chars = DEFAULT_MAX_OUTPUT_CHARS;
    l->max_file_bytes = DEFAULT_MAX_FILE_BYTES;
    sl_init(&l->shell_allowlist);
    sl_init(&l->shell_denylist);
    l->confirm_destructive = false;
    l->verbose = false;
}

void tool_limits_set(const tool_limits *l) {
    limits_ensure();
    tool_limits_free(&g_limits);
    g_limits = *l;
    sl_init(&g_limits.shell_allowlist);
    sl_init(&g_limits.shell_denylist);
    for (size_t i = 0; i < l->shell_allowlist.len; i++)
        sl_push(&g_limits.shell_allowlist, xstrdup(l->shell_allowlist.items[i]));
    for (size_t i = 0; i < l->shell_denylist.len; i++)
        sl_push(&g_limits.shell_denylist, xstrdup(l->shell_denylist.items[i]));
}

const tool_limits *tool_limits_current(void) {
    limits_ensure();
    return &g_limits;
}

void tool_limits_get(tool_limits *out) {
    limits_ensure();
    *out = g_limits;
    sl_init(&out->shell_allowlist);
    sl_init(&out->shell_denylist);
    for (size_t i = 0; i < g_limits.shell_allowlist.len; i++)
        sl_push(&out->shell_allowlist,
                xstrdup(g_limits.shell_allowlist.items[i]));
    for (size_t i = 0; i < g_limits.shell_denylist.len; i++)
        sl_push(&out->shell_denylist, xstrdup(g_limits.shell_denylist.items[i]));
}

void tool_limits_free(tool_limits *l) {
    sl_free(&l->shell_allowlist);
    sl_free(&l->shell_denylist);
}

static bool entry_matches(const char *entry_raw, const char *command_raw) {
    size_t elen, clen;
    const char *entry = xtrim_span(entry_raw, strlen(entry_raw), &elen);
    const char *command = xtrim_span(command_raw, strlen(command_raw), &clen);
    if (elen == 0)
        return true;
    if (clen == elen && strncmp(entry, command, elen) == 0)
        return true;
    /* leading substring on a word boundary */
    if (clen > elen && strncmp(entry, command, elen) == 0) {
        char c = command[elen];
        if (c == ' ' || c == '\t')
            return true;
    }
    /* bare first word */
    size_t fw = 0;
    while (fw < clen && !isspace((unsigned char)command[fw]))
        fw++;
    return fw == elen && strncmp(entry, command, elen) == 0;
}

static bool word_eq(const char *w, size_t n, const char *s);

/* Return the likely executable after common command wrappers. This is a
 * conservative token-level check, not a shell parser. */
static const char *unwrap_command(const char *command) {
    const char *p = command;
    for (int depth = 0; depth < 8; depth++) {
        while (*p && isspace((unsigned char)*p))
            p++;
        const char *word = p;
        while (*p && !isspace((unsigned char)*p))
            p++;
        size_t n = (size_t)(p - word);
        bool wrapper = word_eq(word, n, "sudo") || word_eq(word, n, "env") ||
                       word_eq(word, n, "command") || word_eq(word, n, "nohup") ||
                       word_eq(word, n, "time") || word_eq(word, n, "nice") ||
                       word_eq(word, n, "busybox");
        if (!wrapper)
            return word;
        bool skip_next = false;
        while (*p) {
            while (*p && isspace((unsigned char)*p))
                p++;
            if (!*p)
                return p;
            const char *next = p;
            while (*p && !isspace((unsigned char)*p))
                p++;
            size_t nn = (size_t)(p - next);
            if (skip_next) {
                skip_next = false;
                continue;
            }
            if ((word_eq(word, n, "sudo") || word_eq(word, n, "nice")) &&
                (word_eq(next, nn, "-u") || word_eq(next, nn, "-g") ||
                 word_eq(next, nn, "-n") || word_eq(next, nn, "-C") ||
                 word_eq(next, nn, "--user") || word_eq(next, nn, "--group"))) {
                skip_next = true;
                continue;
            }
            if (next[0] == '-' ||
                ((word_eq(word, n, "env") || word_eq(word, n, "busybox")) &&
                 strchr(next, '=') != NULL))
                continue;
            p = next;
            break;
        }
    }
    while (*p && isspace((unsigned char)*p))
        p++;
    return p;
}

int shell_policy(const char *command, const tool_limits *l, char **reason) {
    for (size_t i = 0; i < l->shell_denylist.len; i++) {
        const char *e = l->shell_denylist.items[i];
        if (entry_matches(e, command) ||
            entry_matches(e, unwrap_command(command))) {
            xbuf b;
            xbuf_init(&b);
            xbuf_fmt(&b, "command matches the denylist entry `%s`", e);
            char *r = xbuf_steal(&b);
            if (reason)
                *reason = r;
            else
                free(r);
            return SHELL_DENIED;
        }
    }
    if (l->shell_allowlist.len > 0) {
        bool any = false;
        for (size_t i = 0; i < l->shell_allowlist.len; i++)
            if (entry_matches(l->shell_allowlist.items[i], command) ||
                entry_matches(l->shell_allowlist.items[i],
                              unwrap_command(command))) {
                any = true;
                break;
            }
        if (!any) {
            xbuf list;
            xbuf_init(&list);
            for (size_t i = 0; i < l->shell_allowlist.len; i++) {
                if (i)
                    xbuf_put(&list, ", ");
                xbuf_put(&list, l->shell_allowlist.items[i]);
            }
            xbuf b;
            xbuf_init(&b);
            xbuf_fmt(&b, "command not in the allowlist (%s)", list.data ? list.data : "");
            xbuf_free(&list);
            char *r = xbuf_steal(&b);
            if (reason)
                *reason = r;
            else
                free(r);
            return SHELL_DENIED;
        }
    }
    return SHELL_ALLOWED;
}

static bool word_eq(const char *w, size_t n, const char *s) {
    size_t sn = strlen(s);
    return n == sn && strncmp(w, s, n) == 0;
}

bool is_destructive(const char *command) {
    /* This is deliberately conservative rather than a shell parser: syntax
     * that can introduce or expand commands requires confirmation. */
    for (const char *p = command; *p; p++) {
        if (*p == ';' || *p == '|' || *p == '&' || *p == '\n' ||
            *p == '`' || *p == '<' || *p == '>' || *p == '$')
            return true;
    }
    /* split the executable and its first argument after common wrappers */
    const char *p = unwrap_command(command);
    while (*p && isspace((unsigned char)*p))
        p++;
    const char *w1 = p;
    while (*p && !isspace((unsigned char)*p))
        p++;
    size_t w1len = (size_t)(p - w1);
    while (*p && isspace((unsigned char)*p))
        p++;
    const char *w2 = p;
    while (*p && !isspace((unsigned char)*p))
        p++;
    size_t w2len = (size_t)(p - w2);

    if (word_eq(w1, w1len, "rm") || word_eq(w1, w1len, "rmdir") ||
        word_eq(w1, w1len, "kill") || word_eq(w1, w1len, "killall") ||
        word_eq(w1, w1len, "pkill") || word_eq(w1, w1len, "truncate") ||
        word_eq(w1, w1len, "mv") || word_eq(w1, w1len, "chmod") ||
        word_eq(w1, w1len, "chown") || word_eq(w1, w1len, "mkfs") ||
        word_eq(w1, w1len, "fdisk") || word_eq(w1, w1len, "parted") ||
        word_eq(w1, w1len, "mount") || word_eq(w1, w1len, "umount") ||
        word_eq(w1, w1len, "swapon") || word_eq(w1, w1len, "swapoff") ||
        word_eq(w1, w1len, "shutdown") || word_eq(w1, w1len, "reboot") ||
        word_eq(w1, w1len, "poweroff") || word_eq(w1, w1len, "halt"))
        return true;
    if (word_eq(w1, w1len, "dd"))
        return strstr(command, "of=") != NULL;
    if (word_eq(w1, w1len, "git"))
        return word_eq(w2, w2len, "reset") || word_eq(w2, w2len, "clean") ||
               word_eq(w2, w2len, "checkout") || word_eq(w2, w2len, "rebase") ||
               word_eq(w2, w2len, "push");
    return false;
}

void set_confirm_hook(bool (*hook)(const char *, void *), void *ud) {
    g_confirm_hook = hook;
    g_confirm_ud = ud;
}

static bool confirm_command(const char *command) {
    if (g_confirm_hook)
        return g_confirm_hook(command, g_confirm_ud);
    return false;
}

int shell_gate(const char *command, const tool_limits *l, char **err) {
    char *reason = NULL;
    if (shell_policy(command, l, &reason) == SHELL_DENIED) {
        xbuf b;
        xbuf_init(&b);
        xbuf_fmt(&b, "run_shell blocked: %s", reason);
        free(reason);
        if (err)
            *err = xbuf_steal(&b);
        return -1;
    }
    if (l->confirm_destructive && is_destructive(command) &&
        !confirm_command(command)) {
        xbuf b;
        xbuf_init(&b);
        xbuf_fmt(&b, "user declined to run destructive command: %s", command);
        char *r = xbuf_steal(&b);
        if (err)
            *err = r;
        else
            free(r);
        return -1;
    }
    return 0;
}