#define _XOPEN_SOURCE 600
#include "shell.h"

#include <errno.h>
#include <pty.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/ioctl.h>
#include <sys/wait.h>
#include <unistd.h>

/* Fallback cap for in-memory PTY capture; matches the config
 * `max_output_chars` default. The caller overrides with the effective value. */
#define DEFAULT_SHELL_CAPTURE_CAP 16000UL

/* Read the current terminal size (0x0 when not a TTY). */
static struct winsize current_winsize(void) {
    struct winsize ws = {0, 0, 0, 0};
    (void)ioctl(STDIN_FILENO, TIOCGWINSZ, &ws);
    return ws;
}

/* Reader thread: stream the shell's PTY output to stdout, collapsing CRLF to
 * a single LF (the real terminal's ONLCR will re-expand it exactly once). */
static void *shell_reader(void *arg) {
    interactive_shell *s = (interactive_shell *) arg;
    unsigned char buf[8192];
    unsigned char out[16384];
    size_t out_len = 0;
    bool pending_cr = false;

    for (;;) {
        ssize_t n = read(s->reader_fd, buf, sizeof buf);
        if (n <= 0)
            break;
        for (ssize_t i = 0; i < n; i++) {
            unsigned char b = buf[i];
            if (pending_cr) {
                if (b == '\n') {
                    out[out_len++] = '\n';
                } else {
                    if (out_len + 2 > sizeof out) {
                        write(STDOUT_FILENO, out, out_len);
                        out_len = 0;
                    }
                    out[out_len++] = '\r';
                    out[out_len++] = b;
                }
                pending_cr = false;
            } else if (b == '\r') {
                pending_cr = true;
            } else {
                out[out_len++] = b;
            }
        }
        if (out_len) {
            (void)write(STDOUT_FILENO, out, out_len);
            pthread_mutex_lock(&s->out_mu);
            if (s->outbuf.len < s->capture_cap)
                xbuf_putn(&s->outbuf, (const char *)out, out_len);
            pthread_mutex_unlock(&s->out_mu);
            out_len = 0;
        }
    }
    if (pending_cr)
        out[out_len++] = '\r';
    if (out_len) {
        (void)write(STDOUT_FILENO, out, out_len);
        pthread_mutex_lock(&s->out_mu);
        if (s->outbuf.len < s->capture_cap)
            xbuf_putn(&s->outbuf, (const char *)out, out_len);
        pthread_mutex_unlock(&s->out_mu);
    }
    pthread_mutex_lock(&s->done_mu);
    s->reader_done = true;
    pthread_mutex_unlock(&s->done_mu);
    return NULL;
}

interactive_shell *shell_spawn(void) {
    const char *shell = getenv("SHELL");
    if (!shell || !*shell)
        shell = "bash";

    char *argv[] = {(char *)shell, "-i", NULL};
    struct winsize ws = current_winsize();

    int master = 0;
    pid_t pid = forkpty(&master, NULL, NULL, &ws);
    if (pid < 0)
        return NULL;
    if (pid == 0) {
        /* Child: reset inherited signal handlers, then exec. */
        signal(SIGINT, SIG_DFL);
        signal(SIGQUIT, SIG_DFL);
        signal(SIGTERM, SIG_DFL);
        execvp(shell, argv);
        _exit(127);
    }

    interactive_shell *s = (interactive_shell *) xmalloc(sizeof *s);
    memset(s, 0, sizeof *s);
    s->pid = (int)pid;
    s->master = master;
    s->last_cwd = NULL;

    int reader_fd = dup(master);
    if (reader_fd < 0) {
        close(master);
        xbuf_free(&s->outbuf);
        pthread_mutex_destroy(&s->out_mu);
        free(s);
        return NULL;
    }
    s->reader_fd = reader_fd;
    s->reader_done = false;
    xbuf_init(&s->outbuf);
    pthread_mutex_init(&s->out_mu, NULL);
    /* Cap the in-memory capture at the tool output limit (the config
     * `max_output_chars` default); the caller overrides this with the
     * effective config value. The transcript is trimmed to the same limit
     * when recorded. */
    s->capture_cap = DEFAULT_SHELL_CAPTURE_CAP;
    if (pthread_mutex_init(&s->done_mu, NULL) != 0) {
        close(reader_fd);
        close(master);
        free(s);
        return NULL;
    }
    if (pthread_create(&s->reader_th, NULL, shell_reader, s) != 0) {
        (void)kill((pid_t)pid, SIGKILL);
        (void)kill(-(pid_t)pid, SIGKILL);
        int status = 0;
        (void)waitpid((pid_t)pid, &status, 0);
        close(reader_fd);
        close(master);
        pthread_mutex_destroy(&s->done_mu);
        xbuf_free(&s->outbuf);
        pthread_mutex_destroy(&s->out_mu);
        free(s);
        return NULL;
    }
    return s;
}

int shell_write(const interactive_shell *s, const void *buf, size_t len) {
    if (s->master < 0)
        return -1;
    const unsigned char *p = (const unsigned char* ) buf;
    size_t off = 0;
    while (off < len) {
        ssize_t n = write(s->master, p + off, len - off);
        if (n < 0) {
            if (errno == EINTR)
                continue;
            return -1;
        }
        off += (size_t)n;
    }
    return 0;
}

void shell_update_winsize(const interactive_shell *s) {
    if (s->master < 0)
        return;
    struct winsize ws = current_winsize();
    (void)ioctl(s->master, TIOCSWINSZ, &ws);
}

/* Read a process's cwd from /proc/<pid>/cwd (Linux). Freshly allocated, or
 * NULL when the process is gone / unreadable. */
static char *read_proc_cwd(int pid) {
    if (pid <= 0)
        return NULL;
    char path[64];
    snprintf(path, sizeof path, "/proc/%d/cwd", pid);
    char link[4096];
    ssize_t n = readlink(path, link, sizeof link - 1);
    if (n <= 0)
        return NULL;
    link[n] = '\0';
    return xstrdup(link);
}

char *shell_final_cwd(const interactive_shell *s) {
    char *live = read_proc_cwd(s->pid);
    if (live)
        return live;
    if (s->last_cwd)
        return xstrdup(s->last_cwd);
    return NULL;
}

/* Strip terminal control/escape sequences from `src` so the text is suitable
 * for the conversation context (no colors, cursor moves, OSC prompts, or
 * prompt decorations). Keeps printable text, tabs, and newlines; a lone CR
 * (not part of the already-collapsed CRLF) becomes a newline. */
static char *shell_sanitize(const char *src) {
    xbuf b;
    xbuf_init(&b);
    const unsigned char *p = (const unsigned char *)src;
    while (*p) {
        if (*p == '\x1B') {
            if (p[1] == '[') {
                /* CSI: params + intermediates, terminated by 0x40-0x7E. */
                const unsigned char *q = p + 2;
                while (*q && *q < 0x40) q++;
                if (*q) q++;
                p = q;
            } else if (p[1] == ']') {
                /* OSC: terminated by BEL (0x07) or ST (ESC \). */
                const unsigned char *q = p + 2;
                while (*q && *q != '\x07' &&
                       !(*q == '\x1B' && q[1] == '\\'))
                    q++;
                if (*q == '\x07')
                    q++;
                else if (*q == '\x1B')
                    q += 2;
                p = q;
            } else {
                /* Other two/three-byte escapes (e.g. ESC ( B, ESC 7). */
                p += (p[1] != 0) ? 2 : 1;
            }
            continue;
        }
        if (*p == '\r') {
            xbuf_push(&b, '\n');
            p++;
            continue;
        }
        if (*p < 0x20 && *p != '\t' && *p != '\n') {
            p++;
            continue;
        }
        xbuf_push(&b, (char)*p);
        p++;
    }
    return xbuf_steal(&b);
}

char *shell_captured(interactive_shell *s) {
    pthread_mutex_lock(&s->out_mu);
    const char *raw = s->outbuf.data ? s->outbuf.data : "";
    char *out = shell_sanitize(raw);
    xbuf_reset(&s->outbuf);
    pthread_mutex_unlock(&s->out_mu);
    return out;
}

bool shell_reap_if_dead(interactive_shell *s) {
    if (s->pid == 0)
        return true;
    /* Wait for the PTY reader to see EOF before returning to the chat prompt,
     * so final shell output is not interleaved with it. */
    pthread_mutex_lock(&s->done_mu);
    bool reader_done = s->reader_done;
    pthread_mutex_unlock(&s->done_mu);
    if (!reader_done) {
        /* Track where the shell is while it is still alive: once it has
         * exited, /proc/<pid>/cwd is unreadable even before the reap. */
        char *cwd = read_proc_cwd(s->pid);
        if (cwd) {
            free(s->last_cwd);
            s->last_cwd = cwd;
        }
        return false;
    }
    int status = 0;
    pid_t r = waitpid((pid_t)s->pid, &status, WNOHANG);
    if (r == (pid_t)s->pid) {
        s->pid = 0;
        return true;
    }
    return false;
}

void shell_free(interactive_shell *s) {
    if (!s)
        return;
    if (s->pid != 0) {
        /* Kill the shell (and its process group) so it does not print a
         * newline or prompt fragment while we redraw the chat prompt. */
        (void)kill((pid_t)s->pid, SIGKILL);
        (void)kill(-(pid_t)s->pid, SIGKILL);
        int status = 0;
        (void)waitpid((pid_t)s->pid, &status, 0);
        s->pid = 0;
    }
    if (s->master >= 0) {
        close(s->master);
        s->master = -1;
    }
    (void)pthread_join(s->reader_th, NULL);
    pthread_mutex_destroy(&s->done_mu);
    xbuf_free(&s->outbuf);
    pthread_mutex_destroy(&s->out_mu);
    free(s->last_cwd);
    s->last_cwd = NULL;
    free(s);
}