#include "Python.h"
#include "pycore_token.h"
#include "errcode.h"

#include "lexer_internal.h"
#include "../tokenizer/helpers.h"

#define MAKE_TOKEN(token_type) _PyLexer_token_setup(tok, token, token_type, p_start, p_end)

int
_PyLexer_set_ftstring_expr(struct tok_state* tok, struct token *token, char c) {
    assert(token != NULL);
    assert(c == '}' || c == ':' || c == '!');
    tokenizer_mode *tok_mode = TOK_GET_MODE(tok);

    if (!(tok_mode->in_debug || tok_mode->string_kind == TSTRING) || token->metadata) {
        return 0;
    }
    PyObject *res = NULL;

    // Look for a # character outside of string literals
    int hash_detected = 0;
    int in_string = 0;
    char quote_char = 0;

    for (Py_ssize_t i = 0; i < tok_mode->last_expr_size - tok_mode->last_expr_end; i++) {
        char ch = tok_mode->last_expr_buffer[i];

        // Skip escaped characters
        if (ch == '\\') {
            i++;
            continue;
        }

        // Handle quotes
        if (ch == '"' || ch == '\'') {
            // The following if/else block works becase there is an off number
            // of quotes in STRING tokens and the lexer only ever reaches this
            // function with valid STRING tokens.
            // For example: """hello"""
            // First quote: in_string = 1
            // Second quote: in_string = 0
            // Third quote: in_string = 1
            if (!in_string) {
                in_string = 1;
                quote_char = ch;
            }
            else if (ch == quote_char) {
                in_string = 0;
            }
            continue;
        }

        // Check for # outside strings
        if (ch == '#' && !in_string) {
            hash_detected = 1;
            break;
        }
    }
    // If we found a # character in the expression, we need to handle comments
    if (hash_detected) {
        // Allocate buffer for processed result
        char *result = (char *)PyMem_Malloc((tok_mode->last_expr_size - tok_mode->last_expr_end + 1) * sizeof(char));
        if (!result) {
            return -1;
        }

        Py_ssize_t i = 0;  // Input position
        Py_ssize_t j = 0;  // Output position
        in_string = 0;     // Whether we're in a string
        quote_char = 0;    // Current string quote char

        // Process each character
        while (i < tok_mode->last_expr_size - tok_mode->last_expr_end) {
            char ch = tok_mode->last_expr_buffer[i];

            // Handle string quotes
            if (ch == '"' || ch == '\'') {
                // See comment above to understand this part
                if (!in_string) {
                    in_string = 1;
                    quote_char = ch;
                } else if (ch == quote_char) {
                    in_string = 0;
                }
                result[j++] = ch;
            }
            // Skip comments
            else if (ch == '#' && !in_string) {
                while (i < tok_mode->last_expr_size - tok_mode->last_expr_end &&
                       tok_mode->last_expr_buffer[i] != '\n') {
                    i++;
                }
                if (i < tok_mode->last_expr_size - tok_mode->last_expr_end) {
                    result[j++] = '\n';
                }
            }
            // Copy other chars
            else {
                result[j++] = ch;
            }
            i++;
        }

        result[j] = '\0';  // Null-terminate the result string
        res = PyUnicode_DecodeUTF8(result, j, NULL);
        PyMem_Free(result);
    } else {
        res = PyUnicode_DecodeUTF8(
            tok_mode->last_expr_buffer,
            tok_mode->last_expr_size - tok_mode->last_expr_end,
            NULL
        );
    }

    if (!res) {
        return -1;
    }
    token->metadata = res;
    return 0;
}

int
_PyLexer_update_ftstring_expr(struct tok_state *tok, char cur)
{
    assert(tok->cur != NULL);

    Py_ssize_t size = strlen(tok->cur);
    tokenizer_mode *tok_mode = TOK_GET_MODE(tok);

    switch (cur) {
       case 0:
            if (!tok_mode->last_expr_buffer || tok_mode->last_expr_end >= 0) {
                return 1;
            }
            char *new_buffer = PyMem_Realloc(
                tok_mode->last_expr_buffer,
                tok_mode->last_expr_size + size
            );
            if (new_buffer == NULL) {
                PyMem_Free(tok_mode->last_expr_buffer);
                goto error;
            }
            tok_mode->last_expr_buffer = new_buffer;
            strncpy(tok_mode->last_expr_buffer + tok_mode->last_expr_size, tok->cur, size);
            tok_mode->last_expr_size += size;
            break;
        case '{':
            if (tok_mode->last_expr_buffer != NULL) {
                PyMem_Free(tok_mode->last_expr_buffer);
            }
            tok_mode->last_expr_buffer = PyMem_Malloc(size);
            if (tok_mode->last_expr_buffer == NULL) {
                goto error;
            }
            tok_mode->last_expr_size = size;
            tok_mode->last_expr_end = -1;
            strncpy(tok_mode->last_expr_buffer, tok->cur, size);
            break;
        case '}':
        case '!':
            tok_mode->last_expr_end = strlen(tok->start);
            break;
        case ':':
            if (tok_mode->last_expr_end == -1) {
               tok_mode->last_expr_end = strlen(tok->start);
            }
            break;
        default:
            Py_UNREACHABLE();
    }
    return 1;
error:
    tok->done = E_NOMEM;
    return 0;
}

int
_PyLexer_check_string_prefixes(struct tok_state *tok,
                                             int saw_b, int saw_r, int saw_u,
                                             int saw_f, int saw_t) {
    // Supported: rb, rf, rt (in any order)
    // Unsupported: ub, ur, uf, ut, bf, bt, ft (in any order)

#define RETURN_SYNTAX_ERROR(PREFIX1, PREFIX2)                             \
    do {                                                                  \
        (void)_PyTokenizer_syntaxerror_known_range(                       \
            tok, (int)(tok->start + 1 - tok->line_start),                 \
            (int)(tok->cur - tok->line_start),                            \
            "'" PREFIX1 "' and '" PREFIX2 "' prefixes are incompatible"); \
        return -1;                                                        \
    } while (0)

    if (saw_u && saw_b) {
        RETURN_SYNTAX_ERROR("u", "b");
    }
    if (saw_u && saw_r) {
        RETURN_SYNTAX_ERROR("u", "r");
    }
    if (saw_u && saw_f) {
        RETURN_SYNTAX_ERROR("u", "f");
    }
    if (saw_u && saw_t) {
        RETURN_SYNTAX_ERROR("u", "t");
    }

    if (saw_b && saw_f) {
        RETURN_SYNTAX_ERROR("b", "f");
    }
    if (saw_b && saw_t) {
        RETURN_SYNTAX_ERROR("b", "t");
    }

    if (saw_f && saw_t) {
        RETURN_SYNTAX_ERROR("f", "t");
    }

#undef RETURN_SYNTAX_ERROR

    return 0;
}

int
_PyLexer_scan_fstring_start(struct tok_state *tok, struct token *token, int c)
{
    const char *p_start = NULL;
    const char *p_end = NULL;

    int quote = c;
    int quote_size = 1;             /* 1 or 3 */

    /* Nodes of type STRING, especially multi line strings
       must be handled differently in order to get both
       the starting line number and the column offset right.
       (cf. issue 16806) */
    tok->first_lineno = tok->lineno;
    tok->multi_line_start = tok->line_start;

    /* Find the quote size and start of string */
    int after_quote = tok_nextc(tok);
    if (after_quote == quote) {
        int after_after_quote = tok_nextc(tok);
        if (after_after_quote == quote) {
            quote_size = 3;
        }
        else {
            // TODO: Check this
            tok_backup(tok, after_after_quote);
            tok_backup(tok, after_quote);
        }
    }
    if (after_quote != quote) {
        tok_backup(tok, after_quote);
    }


    p_start = tok->start;
    p_end = tok->cur;
    if (tok->tok_mode_stack_index + 1 >= MAXFSTRINGLEVEL) {
        return MAKE_TOKEN(_PyTokenizer_syntaxerror(tok, "too many nested f-strings or t-strings"));
    }
    tokenizer_mode *the_current_tok = TOK_NEXT_MODE(tok);
    the_current_tok->kind = TOK_FSTRING_MODE;
    the_current_tok->quote = quote;
    the_current_tok->quote_size = quote_size;
    the_current_tok->start = tok->start;
    the_current_tok->multi_line_start = tok->line_start;
    the_current_tok->first_line = tok->lineno;
    the_current_tok->start_offset = -1;
    the_current_tok->multi_line_start_offset = -1;
    the_current_tok->last_expr_buffer = NULL;
    the_current_tok->last_expr_size = 0;
    the_current_tok->last_expr_end = -1;
    the_current_tok->in_format_spec = 0;
    the_current_tok->in_debug = 0;

    enum string_kind_t string_kind = FSTRING;
    switch (*tok->start) {
        case 'T':
        case 't':
            the_current_tok->raw = Py_TOLOWER(*(tok->start + 1)) == 'r';
            string_kind = TSTRING;
            break;
        case 'F':
        case 'f':
            the_current_tok->raw = Py_TOLOWER(*(tok->start + 1)) == 'r';
            break;
        case 'R':
        case 'r':
            the_current_tok->raw = 1;
            if (Py_TOLOWER(*(tok->start + 1)) == 't') {
                string_kind = TSTRING;
            }
            break;
        default:
            Py_UNREACHABLE();
    }

    the_current_tok->string_kind = string_kind;
    the_current_tok->curly_bracket_depth = 0;
    the_current_tok->curly_bracket_expr_start_depth = -1;
    return string_kind == TSTRING ? MAKE_TOKEN(TSTRING_START) : MAKE_TOKEN(FSTRING_START);
}

int
_PyLexer_scan_string(struct tok_state *tok, struct token *token, int c)
{
    const char *p_start = NULL;
    const char *p_end = NULL;

    int quote = c;
    int quote_size = 1;             /* 1 or 3 */
    int end_quote_size = 0;
    int has_escaped_quote = 0;

    /* Nodes of type STRING, especially multi line strings
       must be handled differently in order to get both
       the starting line number and the column offset right.
       (cf. issue 16806) */
    tok->first_lineno = tok->lineno;
    tok->multi_line_start = tok->line_start;

    /* Find the quote size and start of string */
    c = tok_nextc(tok);
    if (c == quote) {
        c = tok_nextc(tok);
        if (c == quote) {
            quote_size = 3;
        }
        else {
            end_quote_size = 1;     /* empty string found */
        }
    }
    if (c != quote) {
        tok_backup(tok, c);
    }

    /* Get rest of string */
    while (end_quote_size != quote_size) {
        c = tok_nextc(tok);
        if (tok->done == E_ERROR) {
            return MAKE_TOKEN(ERRORTOKEN);
        }
        if (tok->done == E_DECODE) {
            break;
        }
        if (c == EOF || (quote_size == 1 && c == '\n')) {
            assert(tok->multi_line_start != NULL);
            // shift the tok_state's location into
            // the start of string, and report the error
            // from the initial quote character
            tok->cur = (char *)tok->start;
            tok->cur++;
            tok->line_start = tok->multi_line_start;
            int start = tok->lineno;
            tok->lineno = tok->first_lineno;

            if (INSIDE_FSTRING(tok)) {
                /* When we are in an f-string, before raising the
                 * unterminated string literal error, check whether
                 * does the initial quote matches with f-strings quotes
                 * and if it is, then this must be a missing '}' token
                 * so raise the proper error */
                tokenizer_mode *the_current_tok = TOK_GET_MODE(tok);
                if (the_current_tok->quote == quote &&
                    the_current_tok->quote_size == quote_size) {
                    return MAKE_TOKEN(_PyTokenizer_syntaxerror(tok,
                        "%c-string: expecting '}'", TOK_GET_STRING_PREFIX(tok)));
                }
            }

            if (quote_size == 3) {
                _PyTokenizer_syntaxerror(tok, "unterminated triple-quoted string literal"
                                 " (detected at line %d)", start);
                if (c != '\n') {
                    tok->done = E_EOFS;
                }
                return MAKE_TOKEN(ERRORTOKEN);
            }
            else {
                if (has_escaped_quote) {
                    _PyTokenizer_syntaxerror(
                        tok,
                        "unterminated string literal (detected at line %d); "
                        "perhaps you escaped the end quote?",
                        start
                    );
                } else {
                    _PyTokenizer_syntaxerror(
                        tok, "unterminated string literal (detected at line %d)", start
                    );
                }
                if (c != '\n') {
                    tok->done = E_EOLS;
                }
                return MAKE_TOKEN(ERRORTOKEN);
            }
        }
        if (c == quote) {
            end_quote_size += 1;
        }
        else {
            end_quote_size = 0;
            if (c == '\\') {
                c = tok_nextc(tok);  /* skip escaped char */
                if (c == quote) {  /* but record whether the escaped char was a quote */
                    has_escaped_quote = 1;
                }
                if (c == '\r') {
                    c = tok_nextc(tok);
                }
            }
        }
    }

    p_start = tok->start;
    p_end = tok->cur;
    return MAKE_TOKEN(STRING);
}

int
_PyLexer_get_fstring_mode(struct tok_state *tok, tokenizer_mode* current_tok, struct token *token)
{
    const char *p_start = NULL;
    const char *p_end = NULL;
    int end_quote_size = 0;
    int unicode_escape = 0;

    tok->start = tok->cur;
    tok->first_lineno = tok->lineno;
    tok->starting_col_offset = tok->col_offset;

    // If we start with a bracket, we defer to the normal mode as there is nothing for us to tokenize
    // before it.
    int start_char = tok_nextc(tok);
    if (start_char == '{') {
        int peek1 = tok_nextc(tok);
        tok_backup(tok, peek1);
        tok_backup(tok, start_char);
        if (peek1 != '{') {
            current_tok->curly_bracket_expr_start_depth++;
            if (current_tok->curly_bracket_expr_start_depth >= MAX_EXPR_NESTING) {
                return MAKE_TOKEN(_PyTokenizer_syntaxerror(tok,
                    "%c-string: expressions nested too deeply", TOK_GET_STRING_PREFIX(tok)));
            }
            TOK_GET_MODE(tok)->kind = TOK_REGULAR_MODE;
            return _PyLexer_get_normal_mode(tok, current_tok, token);
        }
    }
    else {
        tok_backup(tok, start_char);
    }

    // Check if we are at the end of the string
    for (int i = 0; i < current_tok->quote_size; i++) {
        int quote = tok_nextc(tok);
        if (quote != current_tok->quote) {
            tok_backup(tok, quote);
            goto f_string_middle;
        }
    }

    if (current_tok->last_expr_buffer != NULL) {
        PyMem_Free(current_tok->last_expr_buffer);
        current_tok->last_expr_buffer = NULL;
        current_tok->last_expr_size = 0;
        current_tok->last_expr_end = -1;
    }

    p_start = tok->start;
    p_end = tok->cur;
    tok->tok_mode_stack_index--;
    return MAKE_TOKEN(FTSTRING_END(current_tok));

f_string_middle:

    // TODO: This is a bit of a hack, but it works for now. We need to find a better way to handle
    // this.
    tok->multi_line_start = tok->line_start;
    while (end_quote_size != current_tok->quote_size) {
        int c = tok_nextc(tok);
        if (tok->done == E_ERROR || tok->done == E_DECODE) {
            return MAKE_TOKEN(ERRORTOKEN);
        }
        int in_format_spec = (
                current_tok->in_format_spec
                &&
                INSIDE_FSTRING_EXPR(current_tok)
        );

       if (c == EOF || (current_tok->quote_size == 1 && c == '\n')) {
            if (tok->decoding_erred) {
                return MAKE_TOKEN(ERRORTOKEN);
            }

            // If we are in a format spec and we found a newline,
            // it means that the format spec ends here and we should
            // return to the regular mode.
            if (in_format_spec && c == '\n') {
                if (current_tok->quote_size == 1) {
                    return MAKE_TOKEN(
                        _PyTokenizer_syntaxerror(
                            tok,
                            "%c-string: newlines are not allowed in format specifiers for single quoted %c-strings",
                            TOK_GET_STRING_PREFIX(tok), TOK_GET_STRING_PREFIX(tok)
                        )
                    );
                }
                tok_backup(tok, c);
                TOK_GET_MODE(tok)->kind = TOK_REGULAR_MODE;
                current_tok->in_format_spec = 0;
                p_start = tok->start;
                p_end = tok->cur;
                return MAKE_TOKEN(FTSTRING_MIDDLE(current_tok));
            }

            assert(tok->multi_line_start != NULL);
            // shift the tok_state's location into
            // the start of string, and report the error
            // from the initial quote character
            tok->cur = (char *)current_tok->start;
            tok->cur++;
            tok->line_start = current_tok->multi_line_start;
            int start = tok->lineno;

            tokenizer_mode *the_current_tok = TOK_GET_MODE(tok);
            tok->lineno = the_current_tok->first_line;

            if (current_tok->quote_size == 3) {
                _PyTokenizer_syntaxerror(tok,
                                    "unterminated triple-quoted %c-string literal"
                                    " (detected at line %d)",
                                    TOK_GET_STRING_PREFIX(tok), start);
                if (c != '\n') {
                    tok->done = E_EOFS;
                }
                return MAKE_TOKEN(ERRORTOKEN);
            }
            else {
                return MAKE_TOKEN(_PyTokenizer_syntaxerror(tok,
                                    "unterminated %c-string literal (detected at"
                                    " line %d)", TOK_GET_STRING_PREFIX(tok), start));
            }
        }

        if (c == current_tok->quote) {
            end_quote_size += 1;
            continue;
        } else {
            end_quote_size = 0;
        }

        if (c == '{') {
            if (!_PyLexer_update_ftstring_expr(tok, c)) {
                return MAKE_TOKEN(ENDMARKER);
            }
            int peek = tok_nextc(tok);
            if (peek != '{' || in_format_spec) {
                tok_backup(tok, peek);
                tok_backup(tok, c);
                current_tok->curly_bracket_expr_start_depth++;
                if (current_tok->curly_bracket_expr_start_depth >= MAX_EXPR_NESTING) {
                    return MAKE_TOKEN(_PyTokenizer_syntaxerror(tok,
                        "%c-string: expressions nested too deeply", TOK_GET_STRING_PREFIX(tok)));
                }
                TOK_GET_MODE(tok)->kind = TOK_REGULAR_MODE;
                current_tok->in_format_spec = 0;
                p_start = tok->start;
                p_end = tok->cur;
            } else {
                p_start = tok->start;
                p_end = tok->cur - 1;
            }
            return MAKE_TOKEN(FTSTRING_MIDDLE(current_tok));
        } else if (c == '}') {
            if (unicode_escape) {
                p_start = tok->start;
                p_end = tok->cur;
                return MAKE_TOKEN(FTSTRING_MIDDLE(current_tok));
            }
            int peek = tok_nextc(tok);

            // The tokenizer can only be in the format spec if we have already completed the expression
            // scanning (indicated by the end of the expression being set) and we are not at the top level
            // of the bracket stack (-1 is the top level). Since format specifiers can't legally use double
            // brackets, we can bypass it here.
            int cursor = current_tok->curly_bracket_depth;
            if (peek == '}' && !in_format_spec && cursor == 0) {
                p_start = tok->start;
                p_end = tok->cur - 1;
            } else {
                tok_backup(tok, peek);
                tok_backup(tok, c);
                TOK_GET_MODE(tok)->kind = TOK_REGULAR_MODE;
                current_tok->in_format_spec = 0;
                p_start = tok->start;
                p_end = tok->cur;
            }
            return MAKE_TOKEN(FTSTRING_MIDDLE(current_tok));
        } else if (c == '\\') {
            int peek = tok_nextc(tok);
            if (peek == '\r') {
                peek = tok_nextc(tok);
            }
            // Special case when the backslash is right before a curly
            // brace. We have to restore and return the control back
            // to the loop for the next iteration.
            if (peek == '{' || peek == '}') {
                if (!current_tok->raw) {
                    if (_PyTokenizer_warn_invalid_escape_sequence(tok, peek)) {
                        return MAKE_TOKEN(ERRORTOKEN);
                    }
                }
                tok_backup(tok, peek);
                continue;
            }

            if (!current_tok->raw) {
                if (peek == 'N') {
                    /* Handle named unicode escapes (\N{BULLET}) */
                    peek = tok_nextc(tok);
                    if (peek == '{') {
                        unicode_escape = 1;
                    } else {
                        tok_backup(tok, peek);
                    }
                }
            } /* else {
                skip the escaped character
            }*/
        }
    }

    // Backup the f-string quotes to emit a final FSTRING_MIDDLE and
    // add the quotes to the FSTRING_END in the next tokenizer iteration.
    for (int i = 0; i < current_tok->quote_size; i++) {
        tok_backup(tok, current_tok->quote);
    }
    p_start = tok->start;
    p_end = tok->cur;
    return MAKE_TOKEN(FTSTRING_MIDDLE(current_tok));
}