/* SPDX-License-Identifier: GPL-2.0-only */

#include "toml.h"

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

/*
 * Minimal recursive-descent TOML parser. The subset is documented in
 * toml.h. The design priorities are, in order:
 *   1. Reject things we cannot handle, with a useful error.
 *   2. Be small and obviously correct.
 *   3. Speed (a manifest is < 4 KiB; we do not care).
 */

typedef struct {
    const char *src;
    const char *p;
    int line;
    int col;
    char *err; /* owned; non-NULL means a parse error has been recorded */
} parser_t;

/* ------------------------- node construction ----------------------------- */

static ppm_toml_node_t *node_new(ppm_toml_kind_t k, int line, int col) {
    ppm_toml_node_t *n = calloc(1, sizeof(*n));
    if (!n) return NULL;
    n->kind = k;
    n->line = line;
    n->col  = col;
    return n;
}

static int arr_push(ppm_toml_node_t *arr, ppm_toml_node_t *item) {
    if (arr->u.arr.len == arr->u.arr.cap) {
        size_t cap = arr->u.arr.cap ? arr->u.arr.cap * 2 : 4;
        ppm_toml_node_t **p = realloc(arr->u.arr.items, cap * sizeof(*p));
        if (!p) return -1;
        arr->u.arr.items = p;
        arr->u.arr.cap   = cap;
    }
    arr->u.arr.items[arr->u.arr.len++] = item;
    return 0;
}

static int tbl_push(ppm_toml_node_t *tbl, char *key, ppm_toml_node_t *value) {
    if (tbl->u.tbl.len == tbl->u.tbl.cap) {
        size_t cap = tbl->u.tbl.cap ? tbl->u.tbl.cap * 2 : 4;
        ppm_toml_pair_t *p = realloc(tbl->u.tbl.pairs, cap * sizeof(*p));
        if (!p) return -1;
        tbl->u.tbl.pairs = p;
        tbl->u.tbl.cap   = cap;
    }
    tbl->u.tbl.pairs[tbl->u.tbl.len].key   = key;
    tbl->u.tbl.pairs[tbl->u.tbl.len].value = value;
    tbl->u.tbl.len++;
    return 0;
}

void ppm_toml_free(ppm_toml_node_t *node) {
    if (!node) return;
    switch (node->kind) {
    case PPM_TOML_STRING:
        free(node->u.str);
        break;
    case PPM_TOML_ARRAY:
        for (size_t i = 0; i < node->u.arr.len; i++)
            ppm_toml_free(node->u.arr.items[i]);
        free(node->u.arr.items);
        break;
    case PPM_TOML_TABLE:
        for (size_t i = 0; i < node->u.tbl.len; i++) {
            free(node->u.tbl.pairs[i].key);
            ppm_toml_free(node->u.tbl.pairs[i].value);
        }
        free(node->u.tbl.pairs);
        break;
    }
    free(node);
}

/* ------------------------- lookups --------------------------------------- */

static const ppm_toml_pair_t *find_pair(const ppm_toml_node_t *tbl, const char *key) {
    if (!tbl || tbl->kind != PPM_TOML_TABLE || !key) return NULL;
    for (size_t i = 0; i < tbl->u.tbl.len; i++) {
        if (strcmp(tbl->u.tbl.pairs[i].key, key) == 0)
            return &tbl->u.tbl.pairs[i];
    }
    return NULL;
}

const ppm_toml_node_t *ppm_toml_get(const ppm_toml_node_t *tbl, const char *key) {
    const ppm_toml_pair_t *p = find_pair(tbl, key);
    return p ? p->value : NULL;
}

const char *ppm_toml_get_str(const ppm_toml_node_t *tbl, const char *key) {
    const ppm_toml_node_t *n = ppm_toml_get(tbl, key);
    return (n && n->kind == PPM_TOML_STRING) ? n->u.str : NULL;
}

const ppm_toml_node_t *ppm_toml_get_table(const ppm_toml_node_t *tbl, const char *key) {
    const ppm_toml_node_t *n = ppm_toml_get(tbl, key);
    return (n && n->kind == PPM_TOML_TABLE) ? n : NULL;
}

const ppm_toml_node_t *ppm_toml_get_array(const ppm_toml_node_t *tbl, const char *key) {
    const ppm_toml_node_t *n = ppm_toml_get(tbl, key);
    return (n && n->kind == PPM_TOML_ARRAY) ? n : NULL;
}

/* ------------------------- lexer helpers --------------------------------- */

static void parser_errf(parser_t *p, const char *fmt, ...)
    __attribute__((format(printf, 2, 3)));

static void parser_errf(parser_t *p, const char *fmt, ...) {
    if (p->err) return; /* keep the first error */
    char body[256];
    va_list ap;
    va_start(ap, fmt);
    vsnprintf(body, sizeof(body), fmt, ap);
    va_end(ap);
    char *msg = NULL;
    if (asprintf(&msg, "line %d, col %d: %s", p->line, p->col, body) < 0)
        msg = NULL;
    p->err = msg;
}

static int peek(const parser_t *p) {
    return (unsigned char)*p->p;
}

static int advance(parser_t *p) {
    int c = (unsigned char)*p->p;
    if (c == 0) return 0;
    p->p++;
    if (c == '\n') { p->line++; p->col = 1; }
    else           { p->col++; }
    return c;
}

static void skip_ws_and_comments(parser_t *p) {
    for (;;) {
        int c = peek(p);
        if (c == ' ' || c == '\t' || c == '\r' || c == '\n') { advance(p); continue; }
        if (c == '#') {
            while (peek(p) && peek(p) != '\n') advance(p);
            continue;
        }
        break;
    }
}

static void skip_inline_ws(parser_t *p) {
    while (peek(p) == ' ' || peek(p) == '\t') advance(p);
}

static int is_bare_key_char(int c) {
    return isalnum(c) || c == '-' || c == '_';
}

/* ------------------------- value parsers --------------------------------- */

static char *parse_basic_string(parser_t *p) {
    /* assumes current char is '"' */
    advance(p);
    size_t cap = 32, len = 0;
    char *buf = malloc(cap);
    if (!buf) { parser_errf(p, "out of memory"); return NULL; }

    for (;;) {
        int c = peek(p);
        if (c == 0 || c == '\n') {
            parser_errf(p, "unterminated basic string");
            free(buf);
            return NULL;
        }
        if (c == '"') { advance(p); break; }
        if (c == '\\') {
            advance(p);
            int e = advance(p);
            switch (e) {
            case '"':  c = '"'; break;
            case '\\': c = '\\'; break;
            case 'n':  c = '\n'; break;
            case 't':  c = '\t'; break;
            case 'r':  c = '\r'; break;
            case '0':  c = '\0'; break;
            default:
                parser_errf(p, "unsupported escape sequence \\%c", e);
                free(buf);
                return NULL;
            }
        } else {
            advance(p);
        }
        if (len + 1 >= cap) {
            cap *= 2;
            char *nb = realloc(buf, cap);
            if (!nb) { parser_errf(p, "out of memory"); free(buf); return NULL; }
            buf = nb;
        }
        buf[len++] = (char)c;
    }
    buf[len] = '\0';
    return buf;
}

static char *parse_literal_string(parser_t *p) {
    /* assumes current char is '\'' */
    advance(p);
    const char *start = p->p;
    while (peek(p) && peek(p) != '\'' && peek(p) != '\n') advance(p);
    if (peek(p) != '\'') {
        parser_errf(p, "unterminated literal string");
        return NULL;
    }
    size_t n = (size_t)(p->p - start);
    char *s = strndup(start, n);
    if (!s) { parser_errf(p, "out of memory"); return NULL; }
    advance(p); /* consume closing quote */
    return s;
}

/*
 * Multi-line basic string: """...""". TOML rules:
 *   - opening """ is followed by an optional newline that is stripped
 *   - the body may span lines; embedded "" or " sequences are fine as long
 *     as they are not followed by a third "
 *   - the same escape sequences as basic strings work
 *   - a trailing backslash immediately before a newline trims the newline
 *     and any leading whitespace on the next line (line continuation)
 */
static char *parse_multiline_basic_string(parser_t *p) {
    /* Consume the opening """. The caller has already peeked one '"'. */
    advance(p); advance(p); advance(p);

    /* Strip immediate leading newline if present. */
    if (peek(p) == '\r') advance(p);
    if (peek(p) == '\n') advance(p);

    size_t cap = 64, len = 0;
    char *buf = malloc(cap);
    if (!buf) { parser_errf(p, "out of memory"); return NULL; }

    for (;;) {
        if (peek(p) == 0) {
            parser_errf(p, "unterminated multi-line string");
            free(buf);
            return NULL;
        }

        /* Closing delimiter: three quotes in a row. TOML allows up to
         * two trailing quotes inside the body, so we look ahead. */
        if (peek(p) == '"') {
            int q = 0;
            while (p->p[q] == '"' && q < 5) q++;
            if (q >= 3) {
                /* Up to 5 quotes: 3 close + at most 2 trailing in body. */
                int trailing = (q > 3) ? q - 3 : 0;
                for (int i = 0; i < 3 + trailing; i++) advance(p);
                while (trailing > 0) {
                    if (len + 1 >= cap) {
                        cap *= 2;
                        char *nb = realloc(buf, cap);
                        if (!nb) { parser_errf(p, "out of memory"); free(buf); return NULL; }
                        buf = nb;
                    }
                    buf[len++] = '"';
                    trailing--;
                }
                buf[len] = '\0';
                return buf;
            }
        }

        int c = peek(p);
        if (c == '\\') {
            advance(p);
            int e = peek(p);
            /* Line continuation: backslash + (optional whitespace) + newline
             * + leading whitespace on the next line is all elided. */
            if (e == '\n' || e == '\r' || e == ' ' || e == '\t') {
                /* skip whitespace up to and including the newline */
                while (peek(p) == ' ' || peek(p) == '\t') advance(p);
                if (peek(p) == '\r') advance(p);
                if (peek(p) == '\n') {
                    advance(p);
                    /* eat leading whitespace on subsequent lines */
                    while (peek(p) == ' ' || peek(p) == '\t') advance(p);
                    continue;
                }
                /* not a line continuation; emit literal backslash + whatever */
                c = '\\';
            } else {
                advance(p);
                switch (e) {
                case '"':  c = '"'; break;
                case '\\': c = '\\'; break;
                case 'n':  c = '\n'; break;
                case 't':  c = '\t'; break;
                case 'r':  c = '\r'; break;
                case '0':  c = '\0'; break;
                default:
                    parser_errf(p, "unsupported escape sequence \\%c", e);
                    free(buf);
                    return NULL;
                }
            }
        } else {
            advance(p);
        }

        if (len + 1 >= cap) {
            cap *= 2;
            char *nb = realloc(buf, cap);
            if (!nb) { parser_errf(p, "out of memory"); free(buf); return NULL; }
            buf = nb;
        }
        buf[len++] = (char)c;
    }
}

static char *parse_string_value(parser_t *p) {
    if (peek(p) == '"') {
        /* Detect multi-line """...""" before single-line "...". */
        if (p->p[1] == '"' && p->p[2] == '"') {
            return parse_multiline_basic_string(p);
        }
        return parse_basic_string(p);
    }
    if (peek(p) == '\'') return parse_literal_string(p);
    parser_errf(p, "expected a string");
    return NULL;
}

static ppm_toml_node_t *parse_value(parser_t *p); /* fwd */

static ppm_toml_node_t *parse_array(parser_t *p) {
    int line = p->line, col = p->col;
    advance(p); /* '[' */
    ppm_toml_node_t *arr = node_new(PPM_TOML_ARRAY, line, col);
    if (!arr) { parser_errf(p, "out of memory"); return NULL; }

    for (;;) {
        skip_ws_and_comments(p);
        if (peek(p) == ']') { advance(p); return arr; }

        ppm_toml_node_t *item = parse_value(p);
        if (!item) { ppm_toml_free(arr); return NULL; }

        if (arr->u.arr.len > 0 &&
            arr->u.arr.items[0]->kind != item->kind) {
            parser_errf(p, "heterogeneous arrays are not supported");
            ppm_toml_free(item);
            ppm_toml_free(arr);
            return NULL;
        }
        if (item->kind != PPM_TOML_STRING) {
            parser_errf(p, "only arrays of strings are supported");
            ppm_toml_free(item);
            ppm_toml_free(arr);
            return NULL;
        }

        if (arr_push(arr, item) < 0) {
            parser_errf(p, "out of memory");
            ppm_toml_free(item);
            ppm_toml_free(arr);
            return NULL;
        }

        skip_ws_and_comments(p);
        if (peek(p) == ',') { advance(p); continue; }
        if (peek(p) == ']') { advance(p); return arr; }

        parser_errf(p, "expected ',' or ']' in array");
        ppm_toml_free(arr);
        return NULL;
    }
}

static ppm_toml_node_t *parse_value(parser_t *p) {
    skip_inline_ws(p);
    int line = p->line, col = p->col;
    int c = peek(p);

    if (c == '"' || c == '\'') {
        char *s = parse_string_value(p);
        if (!s) return NULL;
        ppm_toml_node_t *n = node_new(PPM_TOML_STRING, line, col);
        if (!n) { free(s); parser_errf(p, "out of memory"); return NULL; }
        n->u.str = s;
        return n;
    }
    if (c == '[') return parse_array(p);

    parser_errf(p, "unsupported value (only strings and arrays of strings are accepted)");
    return NULL;
}

/* ------------------------- keys & tables --------------------------------- */

static char *parse_key(parser_t *p) {
    skip_inline_ws(p);
    int c = peek(p);
    if (c == '"' || c == '\'') {
        return c == '"' ? parse_basic_string(p) : parse_literal_string(p);
    }
    const char *start = p->p;
    while (is_bare_key_char(peek(p))) advance(p);
    if (p->p == start) {
        parser_errf(p, "expected a key");
        return NULL;
    }
    return strndup(start, (size_t)(p->p - start));
}

/*
 * Walk or create a chain of nested tables for a [section.subsection] header.
 * Returns the final table. The caller must not free it; ownership stays
 * with the document root.
 */
static ppm_toml_node_t *
descend_or_create(parser_t *p, ppm_toml_node_t *root, char **parts, size_t n) {
    ppm_toml_node_t *cur = root;
    for (size_t i = 0; i < n; i++) {
        ppm_toml_node_t *child = NULL;
        for (size_t j = 0; j < cur->u.tbl.len; j++) {
            if (strcmp(cur->u.tbl.pairs[j].key, parts[i]) == 0) {
                child = cur->u.tbl.pairs[j].value;
                break;
            }
        }
        if (child) {
            if (child->kind != PPM_TOML_TABLE) {
                parser_errf(p, "key '%s' already defined as a non-table", parts[i]);
                return NULL;
            }
        } else {
            child = node_new(PPM_TOML_TABLE, p->line, p->col);
            if (!child) { parser_errf(p, "out of memory"); return NULL; }
            char *k = strdup(parts[i]);
            if (!k || tbl_push(cur, k, child) < 0) {
                free(k);
                ppm_toml_free(child);
                parser_errf(p, "out of memory");
                return NULL;
            }
        }
        cur = child;
    }
    return cur;
}

static int parse_table_header(parser_t *p, ppm_toml_node_t *root,
                              ppm_toml_node_t **out_current) {
    advance(p); /* '[' */
    char *parts[16];
    size_t n = 0;

    for (;;) {
        if (n >= 16) {
            parser_errf(p, "table nesting too deep");
            goto fail;
        }
        char *k = parse_key(p);
        if (!k) goto fail;
        parts[n++] = k;
        skip_inline_ws(p);
        if (peek(p) == '.') { advance(p); continue; }
        if (peek(p) == ']') { advance(p); break; }
        parser_errf(p, "expected '.' or ']' in table header");
        goto fail;
    }

    skip_inline_ws(p);
    if (peek(p) && peek(p) != '\n' && peek(p) != '#' && peek(p) != '\r') {
        parser_errf(p, "trailing characters after table header");
        goto fail;
    }

    ppm_toml_node_t *t = descend_or_create(p, root, parts, n);
    if (!t) goto fail;

    for (size_t i = 0; i < n; i++) free(parts[i]);
    *out_current = t;
    return 0;

fail:
    for (size_t i = 0; i < n; i++) free(parts[i]);
    return -1;
}

static int parse_keyval(parser_t *p, ppm_toml_node_t *cur) {
    char *key = parse_key(p);
    if (!key) return -1;

    /* Reject duplicate keys in the current table. */
    for (size_t i = 0; i < cur->u.tbl.len; i++) {
        if (strcmp(cur->u.tbl.pairs[i].key, key) == 0) {
            parser_errf(p, "duplicate key '%s'", key);
            free(key);
            return -1;
        }
    }

    skip_inline_ws(p);
    if (peek(p) != '=') {
        parser_errf(p, "expected '=' after key '%s'", key);
        free(key);
        return -1;
    }
    advance(p);

    ppm_toml_node_t *value = parse_value(p);
    if (!value) { free(key); return -1; }

    if (tbl_push(cur, key, value) < 0) {
        parser_errf(p, "out of memory");
        free(key);
        ppm_toml_free(value);
        return -1;
    }

    skip_inline_ws(p);
    if (peek(p) == '#') while (peek(p) && peek(p) != '\n') advance(p);
    if (peek(p) && peek(p) != '\n' && peek(p) != '\r' && peek(p) != 0) {
        parser_errf(p, "trailing characters after value");
        return -1;
    }
    return 0;
}

/* ------------------------- top-level ------------------------------------- */

ppm_error_t ppm_toml_parse(const char *src, ppm_toml_node_t **out_node,
                           char **out_err) {
    if (!src || !out_node) return PPM_ERR_INVALID;

    parser_t p = { .src = src, .p = src, .line = 1, .col = 1, .err = NULL };
    ppm_toml_node_t *root = node_new(PPM_TOML_TABLE, 1, 1);
    if (!root) {
        if (out_err) *out_err = strdup("out of memory");
        return PPM_ERR_NOMEM;
    }
    ppm_toml_node_t *cur = root;

    for (;;) {
        skip_ws_and_comments(&p);
        if (peek(&p) == 0) break;

        if (peek(&p) == '[') {
            if (parse_table_header(&p, root, &cur) < 0) goto fail;
        } else {
            if (parse_keyval(&p, cur) < 0) goto fail;
        }
    }

    *out_node = root;
    if (out_err) *out_err = NULL;
    return PPM_OK;

fail:
    ppm_toml_free(root);
    if (out_err) *out_err = p.err; else free(p.err);
    return PPM_ERR_PARSE;
}
