#include "session.h"

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

#include <sodium.h>

#define LOG_MODULE "session-prompt"
#define LOG_ENABLE_DBG 0
#include "char32.h"
#include "extract.h"
#include "grid.h"
#include "log.h"
#include "render.h"
#include "search.h"
#include "session-crypto.h"
#include "terminal.h"
#include "vt.h"
#include "xmalloc.h"

#define SCROLLBACK_MAX_BYTES (1u << 20)  /* 1 MB */

static char *
buf_to_utf8(const struct terminal *term)
{
    if (term->search.len == 0)
        return NULL;
    /* Make a NUL-terminated char32_t copy for ac32tombs */
    char32_t *tmp = calloc(term->search.len + 1, sizeof(char32_t));
    if (tmp == NULL)
        return NULL;
    memcpy(tmp, term->search.buf, term->search.len * sizeof(char32_t));
    tmp[term->search.len] = U'\0';
    char *utf8 = ac32tombs(tmp);
    free(tmp);
    return utf8;
}

void
session_picker_init(struct terminal *term)
{
    session_picker_free(term);
    term->session_picker.all = session_list(&term->session_picker.all_count);
    term->session_picker.filtered_cap = term->session_picker.all_count;
    term->session_picker.filtered = term->session_picker.filtered_cap > 0
        ? xmalloc(sizeof(size_t) * term->session_picker.filtered_cap)
        : NULL;
    session_picker_refilter(term);
}

void
session_picker_free(struct terminal *term)
{
    session_free_names(term->session_picker.all,
                       term->session_picker.all_count);
    free(term->session_picker.filtered);
    term->session_picker.all = NULL;
    term->session_picker.all_count = 0;
    term->session_picker.filtered = NULL;
    term->session_picker.filtered_count = 0;
    term->session_picker.filtered_cap = 0;
    term->session_picker.sel = 0;
}

void
session_picker_refilter(struct terminal *term)
{
    char *filter = buf_to_utf8(term);
    term->session_picker.filtered_count = 0;
    for (size_t i = 0; i < term->session_picker.all_count; i++) {
        const char *name = term->session_picker.all[i];
        if (filter == NULL || strstr(name, filter) != NULL)
            term->session_picker.filtered[term->session_picker.filtered_count++] = i;
    }
    free(filter);
    if (term->session_picker.sel >= term->session_picker.filtered_count)
        term->session_picker.sel = term->session_picker.filtered_count > 0
            ? term->session_picker.filtered_count - 1
            : 0;
}

void
session_picker_move(struct terminal *term, int delta)
{
    if (term->session_picker.filtered_count == 0)
        return;
    long n = (long)term->session_picker.filtered_count;
    long cur = (long)term->session_picker.sel + delta;
    /* Wrap */
    cur = ((cur % n) + n) % n;
    term->session_picker.sel = (size_t)cur;
}

bool
session_picker_delete_selected(struct terminal *term)
{
    if (term->session_picker.filtered_count == 0)
        return false;
    size_t idx = term->session_picker.filtered[term->session_picker.sel];
    const char *name = term->session_picker.all[idx];
    if (!session_delete(name))
        return false;
    /* Rebuild the list */
    session_picker_init(term);
    return true;
}

static void
dismiss(struct terminal *term)
{
    session_picker_free(term);
    if (term->session_pending_name != NULL)
        sodium_memzero(term->session_pending_name,
                       strlen(term->session_pending_name));
    free(term->session_pending_name);
    term->session_pending_name = NULL;
    term->session_pending_secure = false;
    if (term->session_pending_load_state != NULL) {
        struct session_state *st = term->session_pending_load_state;
        session_state_free(st);
        free(st);
        term->session_pending_load_state = NULL;
    }
    /* Zero any residual password bytes in the input buffer */
    if (term->search.buf != NULL && term->search.len > 0)
        sodium_memzero(term->search.buf,
                       term->search.len * sizeof(term->search.buf[0]));
    term->search.mode = SEARCH_MODE_NORMAL;
    /* search_cancel clears the search bar; reuse it. We deliberately do not
     * leave selection state behind (no commit semantics for sessions). */
    search_cancel(term);
}

static char *
extract_scrollback_text(struct terminal *term, size_t *out_len)
{
    *out_len = 0;
    struct grid *grid = &term->normal;
    if (grid->rows == NULL)
        return NULL;

    const int total = grid->num_rows;
    const int visible = term->rows;
    const int scrollback = total > visible ? total - visible : 0;

    struct extraction_context *ctx = extract_begin(SELECTION_NONE, true);
    if (ctx == NULL)
        return NULL;
    for (int row_no = -scrollback; row_no < visible; row_no++) {
        int idx = ((grid->offset + row_no) % total + total) % total;
        const struct row *row = grid->rows[idx];
        if (row == NULL)
            continue;
        for (int col = 0; col < grid->num_cols; col++) {
            if (!extract_one(term, row, &row->cells[col], col, ctx)) {
                char *junk = NULL;
                size_t jl = 0;
                extract_finish(ctx, &junk, &jl);
                free(junk);
                return NULL;
            }
        }
    }
    char *text = NULL;
    size_t len = 0;
    if (!extract_finish(ctx, &text, &len))
        return NULL;

    /* Trim leading newlines (uninitialized scrollback) */
    size_t start = 0;
    while (start < len && (text[start] == '\n' || text[start] == ' '))
        start++;
    if (start > 0) {
        memmove(text, text + start, len - start + 1);
        len -= start;
    }

    /* Cap at SCROLLBACK_MAX_BYTES, keeping the newest content. */
    if (len > SCROLLBACK_MAX_BYTES) {
        size_t drop = len - SCROLLBACK_MAX_BYTES;
        const char *nl = memchr(text + drop, '\n', len - drop);
        if (nl != NULL)
            drop = (size_t)(nl + 1 - text);
        memmove(text, text + drop, len - drop + 1);
        len -= drop;
    }

    *out_len = len;
    return text;
}


static void
do_save(struct terminal *term, const char *name)
{
    struct session_state st = {0};
    if (!session_capture(term->ptmx, term->cwd, &st)) {
        LOG_ERR("session_save: could not capture state");
        session_state_free(&st);
        return;
    }
    if (!session_save(name, &st))
        LOG_ERR("session_save: write failed for '%s'", name);
    session_state_free(&st);

    /* Plain save deliberately drops any stale encrypted scrollback so the
     * session reflects the current save. */
    session_delete_enc_scrollback(name);
    session_delete_scrollback(name);
}

static void
do_save_secure(struct terminal *term, const char *name, const char *password)
{
    struct session_state st = {0};
    if (!session_capture(term->ptmx, term->cwd, &st)) {
        LOG_ERR("session_save_secure: could not capture state");
        session_state_free(&st);
        return;
    }
    bool ok = session_save(name, &st);
    session_state_free(&st);
    if (!ok) {
        LOG_ERR("session_save_secure: JSON write failed for '%s'", name);
        return;
    }
    /* Capture scrollback as plain text, encrypt, write sidecar. */
    size_t plain_len = 0;
    char *plain = extract_scrollback_text(term, &plain_len);
    if (plain == NULL || plain_len == 0) {
        /* No scrollback to encrypt — clear any stale ciphertext. */
        free(plain);
        session_delete_enc_scrollback(name);
        session_delete_scrollback(name);
        return;
    }
    unsigned char *blob = NULL;
    size_t blob_len = 0;
    bool enc_ok = session_crypto_encrypt(
        password, (const unsigned char *)plain, plain_len, &blob, &blob_len);
    sodium_memzero(plain, plain_len);
    free(plain);
    if (!enc_ok) {
        LOG_ERR("session_save_secure: encryption failed");
        return;
    }
    if (!session_write_enc_blob(name, blob, blob_len))
        LOG_ERR("session_save_secure: sidecar write failed");
    /* Remove any legacy plain sidecar so we never load it accidentally. */
    session_delete_scrollback(name);
    sodium_memzero(blob, blob_len);
    free(blob);
}

/*
 * Replay a saved scrollback into a freshly created terminal. Walks the bytes
 * once to convert bare LF into CRLF (so each saved line lands at column 0),
 * then runs them through vt_from_slave. Called from the LOAD path right after
 * the new tab is spawned, before any output from the new slave has arrived.
 */
static void
replay_scrollback(struct terminal *new_term, const char *text, size_t len)
{
    if (text == NULL || len == 0 || new_term == NULL)
        return;

    char *buf = xmalloc(len * 2);
    size_t out = 0;
    for (size_t i = 0; i < len; i++) {
        char c = text[i];
        if (c == '\n')
            buf[out++] = '\r';
        buf[out++] = c;
    }
    if (out > 0)
        buf[out++] = '\n';

    /* Visually separate replayed content from the upcoming shell prompt */
    /* (single trailing newline is enough; the shell will emit its own CRLF) */

    vt_from_slave(new_term, (const uint8_t *)buf, out);
    free(buf);
}

void
session_prompt_confirm_overwrite(struct terminal *term)
{
    if (term->session_pending_name == NULL) {
        dismiss(term);
        return;
    }
    if (term->session_pending_secure) {
        /* Don't save yet — transition to password prompt. */
        term->search.mode = SEARCH_MODE_SESSION_SAVE_SECURE_PASSWORD;
        term->search.len = 0;
        term->search.cursor = 0;
        if (term->search.buf != NULL)
            term->search.buf[0] = U'\0';
        render_refresh_search(term);
        return;
    }
    do_save(term, term->session_pending_name);
    dismiss(term);
}

void
session_prompt_cancel(struct terminal *term)
{
    dismiss(term);
}

static void
spawn_and_replay(struct terminal *term, struct session_state *st,
                 const char *name, const char *password)
{
    struct terminal *nt = term_tab_new_with_cwd(
        term, st->cwd, st->argc, st->argv, NULL,
        term->shutdown.cb, term->shutdown.cb_data);
    if (nt == NULL) {
        LOG_ERR("session_load: could not spawn new tab");
        return;
    }
    /* Encrypted sidecar takes priority if present. */
    if (password != NULL && session_has_enc_scrollback(name)) {
        unsigned char *blob = NULL;
        size_t blob_len = 0;
        if (session_read_enc_blob(name, &blob, &blob_len)) {
            unsigned char *plain = NULL;
            size_t plain_len = 0;
            if (session_crypto_decrypt(password, blob, blob_len,
                                       &plain, &plain_len)) {
                replay_scrollback(nt, (const char *)plain, plain_len);
                sodium_memzero(plain, plain_len);
                free(plain);
            } else {
                LOG_ERR("session_load: decryption failed for '%s'", name);
            }
            free(blob);
        }
    } else {
        /* Plain sidecar (legacy or never-encrypted). */
        char *sb_text = NULL;
        size_t sb_len = 0;
        if (session_load_scrollback(name, &sb_text, &sb_len))
            replay_scrollback(nt, sb_text, sb_len);
        free(sb_text);
    }
}

void
session_prompt_commit(struct terminal *term)
{
    enum search_mode mode = term->search.mode;

    if (mode == SEARCH_MODE_SESSION_SAVE ||
        mode == SEARCH_MODE_SESSION_SAVE_SECURE_NAME) {
        const bool secure = (mode == SEARCH_MODE_SESSION_SAVE_SECURE_NAME);
        char *name = buf_to_utf8(term);
        if (name == NULL || !session_name_is_valid(name)) {
            LOG_ERR("session name must be non-empty [A-Za-z0-9._-]");
            free(name);
            dismiss(term);
            return;
        }
        if (session_exists(name)) {
            free(term->session_pending_name);
            term->session_pending_name = name;
            term->session_pending_secure = secure;
            term->search.mode = SEARCH_MODE_SESSION_OVERWRITE_CONFIRM;
            render_refresh_search(term);
            return;
        }
        if (secure) {
            /* No overwrite needed, go straight to password prompt. */
            free(term->session_pending_name);
            term->session_pending_name = name;
            term->session_pending_secure = true;
            term->search.mode = SEARCH_MODE_SESSION_SAVE_SECURE_PASSWORD;
            term->search.len = 0;
            term->search.cursor = 0;
            if (term->search.buf != NULL)
                term->search.buf[0] = U'\0';
            render_refresh_search(term);
            return;
        }
        do_save(term, name);
        free(name);
    } else if (mode == SEARCH_MODE_SESSION_SAVE_SECURE_PASSWORD) {
        char *password = buf_to_utf8(term);
        if (password == NULL || password[0] == '\0') {
            LOG_ERR("session_save_secure: empty password");
            free(password);
            dismiss(term);
            return;
        }
        const char *name = term->session_pending_name;
        if (name == NULL) {
            sodium_memzero(password, strlen(password));
            free(password);
            dismiss(term);
            return;
        }
        do_save_secure(term, name, password);
        sodium_memzero(password, strlen(password));
        free(password);
    } else if (mode == SEARCH_MODE_SESSION_LOAD) {
        if (term->session_picker.filtered_count == 0) {
            LOG_ERR("session_load: no matching session");
            dismiss(term);
            return;
        }
        size_t idx = term->session_picker.filtered[term->session_picker.sel];
        const char *picked = term->session_picker.all[idx];

        struct session_state *st = xmalloc(sizeof(*st));
        memset(st, 0, sizeof(*st));
        if (!session_load(picked, st)) {
            LOG_ERR("session_load: could not read '%s'", picked);
            free(st);
            dismiss(term);
            return;
        }
        if (session_has_enc_scrollback(picked)) {
            /* Stash state, transition to password mode. */
            free(term->session_pending_name);
            term->session_pending_name = xstrdup(picked);
            term->session_pending_load_state = st;
            term->search.mode = SEARCH_MODE_SESSION_LOAD_PASSWORD;
            term->search.len = 0;
            term->search.cursor = 0;
            if (term->search.buf != NULL)
                term->search.buf[0] = U'\0';
            session_picker_free(term);
            render_refresh_search(term);
            return;
        }
        /* No encrypted scrollback — spawn now. */
        spawn_and_replay(term, st, picked, NULL);
        session_state_free(st);
        free(st);
    } else if (mode == SEARCH_MODE_SESSION_LOAD_PASSWORD) {
        char *password = buf_to_utf8(term);
        struct session_state *st = term->session_pending_load_state;
        const char *name = term->session_pending_name;
        if (st != NULL && name != NULL && password != NULL) {
            spawn_and_replay(term, st, name, password);
        }
        if (st != NULL) {
            session_state_free(st);
            free(st);
            term->session_pending_load_state = NULL;
        }
        if (password != NULL) {
            sodium_memzero(password, strlen(password));
            free(password);
        }
    }

    dismiss(term);
}