#include "Python.h"
#include "pycore_call.h"          // _PyObject_CallNoArgs()
#include "pycore_fileutils.h"     // _Py_UniversalNewlineFgetsWithSize()
#include "pycore_runtime.h"       // _Py_ID()

#include "errcode.h"              // E_NOMEM

#ifdef HAVE_UNISTD_H
#  include <unistd.h>             // lseek(), read()
#endif

#include "helpers.h"
#include "../lexer/state.h"
#include "../lexer/lexer.h"
#include "../lexer/buffer.h"


static int
tok_concatenate_interactive_new_line(struct tok_state *tok, const char *line) {
    assert(tok->fp_interactive);

    if (!line) {
        return 0;
    }

    Py_ssize_t current_size = tok->interactive_src_end - tok->interactive_src_start;
    Py_ssize_t line_size = strlen(line);
    char last_char = line[line_size > 0 ? line_size - 1 : line_size];
    if (last_char != '\n') {
        line_size += 1;
    }
    char* new_str = tok->interactive_src_start;

    new_str = PyMem_Realloc(new_str, current_size + line_size + 1);
    if (!new_str) {
        if (tok->interactive_src_start) {
            PyMem_Free(tok->interactive_src_start);
        }
        tok->interactive_src_start = NULL;
        tok->interactive_src_end = NULL;
        tok->done = E_NOMEM;
        return -1;
    }
    strcpy(new_str + current_size, line);
    tok->implicit_newline = 0;
    if (last_char != '\n') {
        /* Last line does not end in \n, fake one */
        new_str[current_size + line_size - 1] = '\n';
        new_str[current_size + line_size] = '\0';
        tok->implicit_newline = 1;
    }
    tok->interactive_src_start = new_str;
    tok->interactive_src_end = new_str + current_size + line_size;
    return 0;
}

static int
tok_readline_raw(struct tok_state *tok)
{
    do {
        if (!_PyLexer_tok_reserve_buf(tok, BUFSIZ)) {
            return 0;
        }
        int n_chars = (int)(tok->end - tok->inp);
        size_t line_size = 0;
        char *line = _Py_UniversalNewlineFgetsWithSize(tok->inp, n_chars, tok->fp, NULL, &line_size);
        if (line == NULL) {
            return 1;
        }
        if (tok->fp_interactive &&
            tok_concatenate_interactive_new_line(tok, line) == -1) {
            return 0;
        }
        tok->inp += line_size;
        if (tok->inp == tok->buf) {
            return 0;
        }
    } while (tok->inp[-1] != '\n');
    return 1;
}

static int
tok_readline_recode(struct tok_state *tok) {
    PyObject *line;
    const  char *buf;
    Py_ssize_t buflen;
    line = tok->decoding_buffer;
    if (line == NULL) {
        line = PyObject_CallNoArgs(tok->decoding_readline);
        if (line == NULL) {
            _PyTokenizer_error_ret(tok);
            goto error;
        }
    }
    else {
        tok->decoding_buffer = NULL;
    }
    buf = PyUnicode_AsUTF8AndSize(line, &buflen);
    if (buf == NULL) {
        _PyTokenizer_error_ret(tok);
        goto error;
    }
    // Make room for the null terminator *and* potentially
    // an extra newline character that we may need to artificially
    // add.
    size_t buffer_size = buflen + 2;
    if (!_PyLexer_tok_reserve_buf(tok, buffer_size)) {
        goto error;
    }
    memcpy(tok->inp, buf, buflen);
    tok->inp += buflen;
    *tok->inp = '\0';
    if (tok->fp_interactive &&
        tok_concatenate_interactive_new_line(tok, buf) == -1) {
        goto error;
    }
    Py_DECREF(line);
    return 1;
error:
    Py_XDECREF(line);
    return 0;
}

/* Fetch the next byte from TOK. */
static int fp_getc(struct tok_state *tok) {
    return getc(tok->fp);
}

/* Unfetch the last byte back into TOK.  */
static void fp_ungetc(int c, struct tok_state *tok) {
    ungetc(c, tok->fp);
}

/* Set the readline function for TOK to a StreamReader's
   readline function. The StreamReader is named ENC.

   This function is called from _PyTokenizer_check_bom and _PyTokenizer_check_coding_spec.

   ENC is usually identical to the future value of tok->encoding,
   except for the (currently unsupported) case of UTF-16.

   Return 1 on success, 0 on failure. */
static int
fp_setreadl(struct tok_state *tok, const char* enc)
{
    PyObject *readline, *open, *stream;
    int fd;
    long pos;

    fd = fileno(tok->fp);
    /* Due to buffering the file offset for fd can be different from the file
     * position of tok->fp.  If tok->fp was opened in text mode on Windows,
     * its file position counts CRLF as one char and can't be directly mapped
     * to the file offset for fd.  Instead we step back one byte and read to
     * the end of line.*/
    pos = ftell(tok->fp);
    if (pos == -1 ||
        lseek(fd, (off_t)(pos > 0 ? pos - 1 : pos), SEEK_SET) == (off_t)-1) {
        PyErr_SetFromErrnoWithFilename(PyExc_OSError, NULL);
        return 0;
    }

    open = PyImport_ImportModuleAttrString("io", "open");
    if (open == NULL) {
        return 0;
    }
    stream = PyObject_CallFunction(open, "isisOOO",
                    fd, "r", -1, enc, Py_None, Py_None, Py_False);
    Py_DECREF(open);
    if (stream == NULL) {
        return 0;
    }

    readline = PyObject_GetAttr(stream, &_Py_ID(readline));
    Py_DECREF(stream);
    if (readline == NULL) {
        return 0;
    }
    Py_XSETREF(tok->decoding_readline, readline);

    if (pos > 0) {
        PyObject *bufobj = _PyObject_CallNoArgs(readline);
        if (bufobj == NULL) {
            return 0;
        }
        Py_DECREF(bufobj);
    }

    return 1;
}

static int
tok_underflow_interactive(struct tok_state *tok) {
    if (tok->interactive_underflow == IUNDERFLOW_STOP) {
        tok->done = E_INTERACT_STOP;
        return 1;
    }
    char *newtok = PyOS_Readline(tok->fp ? tok->fp : stdin, stdout, tok->prompt);
    if (newtok != NULL) {
        char *translated = _PyTokenizer_translate_newlines(newtok, 0, 0, tok);
        PyMem_Free(newtok);
        if (translated == NULL) {
            return 0;
        }
        newtok = translated;
    }
    if (tok->encoding && newtok && *newtok) {
        /* Recode to UTF-8 */
        Py_ssize_t buflen;
        const char* buf;
        PyObject *u = _PyTokenizer_translate_into_utf8(newtok, tok->encoding);
        PyMem_Free(newtok);
        if (u == NULL) {
            tok->done = E_DECODE;
            return 0;
        }
        buflen = PyBytes_GET_SIZE(u);
        buf = PyBytes_AS_STRING(u);
        newtok = PyMem_Malloc(buflen+1);
        if (newtok == NULL) {
            Py_DECREF(u);
            tok->done = E_NOMEM;
            return 0;
        }
        strcpy(newtok, buf);
        Py_DECREF(u);
    }
    if (tok->fp_interactive &&
        tok_concatenate_interactive_new_line(tok, newtok) == -1) {
        PyMem_Free(newtok);
        return 0;
    }
    if (tok->nextprompt != NULL) {
        tok->prompt = tok->nextprompt;
    }
    if (newtok == NULL) {
        tok->done = E_INTR;
    }
    else if (*newtok == '\0') {
        PyMem_Free(newtok);
        tok->done = E_EOF;
    }
    else if (tok->start != NULL) {
        Py_ssize_t cur_multi_line_start = tok->multi_line_start - tok->buf;
        _PyLexer_remember_fstring_buffers(tok);
        size_t size = strlen(newtok);
        ADVANCE_LINENO();
        if (!_PyLexer_tok_reserve_buf(tok, size + 1)) {
            PyMem_Free(tok->buf);
            tok->buf = NULL;
            PyMem_Free(newtok);
            return 0;
        }
        memcpy(tok->cur, newtok, size + 1);
        PyMem_Free(newtok);
        tok->inp += size;
        tok->multi_line_start = tok->buf + cur_multi_line_start;
        _PyLexer_restore_fstring_buffers(tok);
    }
    else {
        _PyLexer_remember_fstring_buffers(tok);
        ADVANCE_LINENO();
        PyMem_Free(tok->buf);
        tok->buf = newtok;
        tok->cur = tok->buf;
        tok->line_start = tok->buf;
        tok->inp = strchr(tok->buf, '\0');
        tok->end = tok->inp + 1;
        _PyLexer_restore_fstring_buffers(tok);
    }
    if (tok->done != E_OK) {
        if (tok->prompt != NULL) {
            PySys_WriteStderr("\n");
        }
        return 0;
    }

    if (tok->tok_mode_stack_index && !_PyLexer_update_ftstring_expr(tok, 0)) {
        return 0;
    }
    return 1;
}

static int
tok_underflow_file(struct tok_state *tok)
{
    if (tok->decoding_state == STATE_INIT) {
        /* We have not yet determined the encoding.
           If an encoding is found, use the file-pointer
           reader functions from now on. */
        if (!_PyTokenizer_check_bom(fp_getc, fp_ungetc, fp_setreadl, tok)) {
            _PyTokenizer_error_ret(tok);
            return 0;
        }
        assert(tok->decoding_state != STATE_INIT);
    }
    int raw = tok->decoding_readline == NULL;
    if (raw && tok->decoding_state != STATE_NORMAL) {
        /* Keep the first line in the buffer to validate it later if
         * the encoding has not yet been determined. */
    }
    else if (tok->start == NULL && !INSIDE_FSTRING(tok)) {
        tok->cur = tok->inp = tok->buf;
    }
    /* Read until '\n' or EOF */
    if (!raw) {
        /* We already have a codec associated with this input. */
        if (!tok_readline_recode(tok)) {
            return 0;
        }
    }
    else {
        /* We want a 'raw' read. */
        if (!tok_readline_raw(tok)) {
            return 0;
        }
    }
    if (tok->inp == tok->cur) {
        tok->done = E_EOF;
        return 0;
    }
    tok->implicit_newline = 0;
    if (tok->inp[-1] != '\n') {
        assert(tok->inp + 1 < tok->end);
        /* Last line does not end in \n, fake one */
        *tok->inp++ = '\n';
        *tok->inp = '\0';
        tok->implicit_newline = 1;
    }

    if (tok->tok_mode_stack_index && !_PyLexer_update_ftstring_expr(tok, 0)) {
        return 0;
    }

    ADVANCE_LINENO();
    if (tok->decoding_state != STATE_NORMAL) {
        if (!_PyTokenizer_check_coding_spec(tok->cur, strlen(tok->cur),
                                    tok, fp_setreadl))
        {
            return 0;
        }
        if (tok->lineno >= 2) {
            tok->decoding_state = STATE_NORMAL;
        }
    }
    if (raw && tok->decoding_state == STATE_NORMAL) {
        const char *line = tok->lineno <= 2 ? tok->buf : tok->cur;
        int lineno = tok->lineno <= 2 ? 1 : tok->lineno;
        if (!tok->encoding) {
            /* The default encoding is UTF-8, so make sure we don't have any
               non-UTF-8 sequences in it. */
            if (!_PyTokenizer_ensure_utf8(line, tok, lineno)) {
                _PyTokenizer_error_ret(tok);
                return 0;
            }
        }
        else {
            PyObject *tmp = PyUnicode_Decode(line, strlen(line),
                                             tok->encoding, NULL);
            if (tmp == NULL) {
                _PyTokenizer_error_ret(tok);
                return 0;
            }
            Py_DECREF(tmp);
        }
    }
    assert(tok->done == E_OK);
    return tok->done == E_OK;
}

/* Set up tokenizer for file */
struct tok_state *
_PyTokenizer_FromFile(FILE *fp, const char* enc,
                      const char *ps1, const char *ps2)
{
    struct tok_state *tok = _PyTokenizer_tok_new();
    if (tok == NULL)
        return NULL;
    if ((tok->buf = (char *)PyMem_Malloc(BUFSIZ)) == NULL) {
        _PyTokenizer_Free(tok);
        return NULL;
    }
    tok->cur = tok->inp = tok->buf;
    tok->end = tok->buf + BUFSIZ;
    tok->fp = fp;
    tok->prompt = ps1;
    tok->nextprompt = ps2;
    if (ps1 || ps2) {
        tok->underflow = &tok_underflow_interactive;
    } else {
        tok->underflow = &tok_underflow_file;
    }
    if (enc != NULL) {
        /* Must copy encoding declaration since it
           gets copied into the parse tree. */
        tok->encoding = _PyTokenizer_new_string(enc, strlen(enc), tok);
        if (!tok->encoding) {
            _PyTokenizer_Free(tok);
            return NULL;
        }
        tok->decoding_state = STATE_NORMAL;
    }
    return tok;
}

#if defined(__wasi__) || (defined(__EMSCRIPTEN__) && (__EMSCRIPTEN_major__ >= 3))
// fdopen() with borrowed fd. WASI does not provide dup() and Emscripten's
// dup() emulation with open() is slow.
typedef union {
    void *cookie;
    int fd;
} borrowed;

static ssize_t
borrow_read(void *cookie, char *buf, size_t size)
{
    borrowed b = {.cookie = cookie};
    return read(b.fd, (void *)buf, size);
}

static FILE *
fdopen_borrow(int fd) {
    // supports only reading. seek fails. close and write are no-ops.
    cookie_io_functions_t io_cb = {borrow_read, NULL, NULL, NULL};
    borrowed b = {.fd = fd};
    return fopencookie(b.cookie, "r", io_cb);
}
#else
static FILE *
fdopen_borrow(int fd) {
    fd = _Py_dup(fd);
    if (fd < 0) {
        return NULL;
    }
    return fdopen(fd, "r");
}
#endif

/* Get the encoding of a Python file. Check for the coding cookie and check if
   the file starts with a BOM.

   _PyTokenizer_FindEncodingFilename() returns NULL when it can't find the
   encoding in the first or second line of the file (in which case the encoding
   should be assumed to be UTF-8).

   The char* returned is malloc'ed via PyMem_Malloc() and thus must be freed
   by the caller. */
char *
_PyTokenizer_FindEncodingFilename(int fd, PyObject *filename)
{
    struct tok_state *tok;
    FILE *fp;
    char *encoding = NULL;

    fp = fdopen_borrow(fd);
    if (fp == NULL) {
        return NULL;
    }
    tok = _PyTokenizer_FromFile(fp, NULL, NULL, NULL);
    if (tok == NULL) {
        fclose(fp);
        return NULL;
    }
    if (filename != NULL) {
        tok->filename = Py_NewRef(filename);
    }
    else {
        tok->filename = PyUnicode_FromString("<string>");
        if (tok->filename == NULL) {
            fclose(fp);
            _PyTokenizer_Free(tok);
            return encoding;
        }
    }
    struct token token;
    // We don't want to report warnings here because it could cause infinite recursion
    // if fetching the encoding shows a warning.
    tok->report_warnings = 0;
    while (tok->lineno < 2 && tok->done == E_OK) {
        _PyToken_Init(&token);
        _PyTokenizer_Get(tok, &token);
        _PyToken_Free(&token);
    }
    fclose(fp);
    if (tok->encoding) {
        encoding = (char *)PyMem_Malloc(strlen(tok->encoding) + 1);
        if (encoding) {
            strcpy(encoding, tok->encoding);
        }
    }
    _PyTokenizer_Free(tok);
    return encoding;
}