#include "markdown_io.h"
#include "strbuf.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <ctype.h>

static int match_numbered_prefix(const char *line, int len, int *prefixLen) {
    int i = 0;
    while (i < len && isdigit((unsigned char)line[i])) i++;
    if (i == 0 || i >= len || line[i] != '.') return 0;
    i++;
    if (i < len && line[i] == ' ') i++;
    else if (i < len) return 0; /* digits+'.' followed by something other than a space: not a list marker */
    *prefixLen = i;
    return 1;
}

static int is_blank_line(const char *line, int len) {
    for (int i = 0; i < len; i++) {
        if (line[i] != ' ' && line[i] != '\t') return 0;
    }
    return 1;
}

static int is_code_fence(const char *line, int len) {
    return len >= 3 && line[0] == '`' && line[1] == '`' && line[2] == '`';
}

/* A line consisting of 3+ of the same char among '-', '*', '_' (spaces allowed between them,
   as in "- - -") and nothing else -- a horizontal rule. */
static int is_hr_line(const char *line, int len) {
    int n = 0;
    char c = 0;
    for (int i = 0; i < len; i++) {
        if (line[i] == ' ' || line[i] == '\t') continue;
        if (c == 0) {
            c = line[i];
            if (c != '-' && c != '*' && c != '_') return 0;
            n++;
        } else if (line[i] == c) {
            n++;
        } else {
            return 0;
        }
    }
    return n >= 3;
}

/* Whole-line image: "![alt](url)" and nothing else (trailing spaces/tabs tolerated, matching
   is_hr_line's tolerance). A non-empty URL is required. Mixed inline images embedded in prose
   are explicitly out of scope -- fall through to BLOCK_PARAGRAPH like today. */
static int is_image_line(const char *line, int len, int *altStart, int *altLen, int *urlStart, int *urlLen) {
    while (len > 0 && (line[len - 1] == ' ' || line[len - 1] == '\t')) len--;
    if (len < 5 || line[0] != '!' || line[1] != '[') return 0;
    int closeBracket = -1;
    for (int i = 2; i < len; i++) { if (line[i] == ']') { closeBracket = i; break; } }
    if (closeBracket < 0 || closeBracket + 1 >= len || line[closeBracket + 1] != '(') return 0;
    if (line[len - 1] != ')') return 0;
    *altStart = 2;
    *altLen = closeBracket - 2;
    *urlStart = closeBracket + 2;
    *urlLen = (len - 1) - *urlStart;
    return *urlLen > 0;
}

#define TABLE_MAX_COLS 64

/* Splits a "|"-delimited row into cells: strips one leading/trailing '|' if present (GFM allows
   both "| a | b |" and "a | b" forms), splits the remainder on '|' (no backslash-escape handling
   -- matches every other matcher in this file), trims surrounding spaces/tabs from each cell.
   Writes up to TABLE_MAX_COLS cells (byte ranges into `line`) into cellStart/cellLen. Returns the
   cell count. */
static int split_table_row(const char *line, int len, int *cellStart, int *cellLen) {
    int start = 0, end = len;
    if (end > start && line[start] == '|') start++;
    if (end > start && line[end - 1] == '|') end--;

    int count = 0;
    int segStart = start;
    for (int i = start; i <= end; i++) {
        if (i < end && line[i] != '|') continue;
        int s = segStart, e = i;
        while (s < e && (line[s] == ' ' || line[s] == '\t')) s++;
        while (e > s && (line[e - 1] == ' ' || line[e - 1] == '\t')) e--;
        if (count < TABLE_MAX_COLS) {
            cellStart[count] = s;
            cellLen[count] = e - s;
            count++;
        }
        segStart = i + 1;
    }
    return count;
}

/* A GFM table delimiter row: "|"-separated cells each matching "^:?-+:?$" (at least one '-',
   optional leading/trailing ':' for alignment). Fills outAligns[i] with 'l'/'c'/'r' per cell
   ('l' is also the default with no colons at all). Returns the cell count via outCols. */
static int is_table_delimiter_row(const char *line, int len, int *outCols, char *outAligns) {
    int cellStart[TABLE_MAX_COLS], cellLen[TABLE_MAX_COLS];
    int cols = split_table_row(line, len, cellStart, cellLen);
    if (cols == 0) return 0;
    for (int i = 0; i < cols; i++) {
        int s = cellStart[i], e = s + cellLen[i];
        if (s >= e) return 0;
        int leftColon = line[s] == ':';
        int rightColon = line[e - 1] == ':';
        int dashStart = s + (leftColon ? 1 : 0);
        int dashEnd = e - (rightColon ? 1 : 0);
        if (dashEnd <= dashStart) return 0;
        for (int j = dashStart; j < dashEnd; j++) { if (line[j] != '-') return 0; }
        outAligns[i] = rightColon ? (leftColon ? 'c' : 'r') : 'l';
    }
    *outCols = cols;
    return 1;
}

/* Reads one line's bounds starting at *pos (same \r-stripping logic as the main load loop), does
   NOT advance *pos -- a pure peek. */
static void peek_line(const char *buf, int len, int pos, int *lineStart, int *lineEnd) {
    int start = pos;
    while (pos < len && buf[pos] != '\n') pos++;
    int end = pos;
    if (end > start && buf[end - 1] == '\r') end--;
    *lineStart = start;
    *lineEnd = end;
}

/* Advances *pos past one line (mirrors the main load loop's own line-consuming step). */
static void consume_line(const char *buf, int len, int *pos) {
    while (*pos < len && buf[*pos] != '\n') (*pos)++;
    if (*pos < len) (*pos)++;
}

/* Tries to parse a GFM table starting with `headerLine` (already isolated, not yet consumed from
   *pos) as its header row. Only commits (consuming lines for real, inserting blocks) if the very
   next line validates as a delimiter row with a matching cell count -- GFM's own disambiguation
   rule, so a line that merely contains '|' but isn't followed by a real delimiter row falls
   through to BLOCK_PARAGRAPH untouched, same as before this function existed. Returns 1 if a
   table was parsed (and *pos advanced past it), 0 otherwise (*pos untouched). */
static int try_parse_table(Document *doc, const char *buf, int len, int *pos, const char *headerLine, int headerLineLen) {
    int headerStart[TABLE_MAX_COLS], headerLen[TABLE_MAX_COLS];
    int cols = split_table_row(headerLine, headerLineLen, headerStart, headerLen);
    if (cols == 0) return 0;

    int delimStart, delimEnd;
    peek_line(buf, len, *pos, &delimStart, &delimEnd);
    char aligns[TABLE_MAX_COLS];
    int delimCols;
    if (!is_table_delimiter_row(buf + delimStart, delimEnd - delimStart, &delimCols, aligns) || delimCols != cols) {
        return 0;
    }
    consume_line(buf, len, pos); /* the delimiter row itself is not stored as a block */

    for (int c = 0; c < cols; c++) {
        document_insert_block(doc, doc->count, BLOCK_TABLE_HEADER_CELL, headerLine + headerStart[c], headerLen[c]);
        Block *b = &doc->blocks[doc->count - 1];
        b->tableCol = c;
        b->tableCols = cols;
        b->tableAlign = aligns[c];
    }

    for (;;) {
        int rowStart, rowEnd;
        peek_line(buf, len, *pos, &rowStart, &rowEnd);
        if (rowStart >= len) break; /* EOF */
        int rowLen = rowEnd - rowStart;
        int isBlank = 1;
        for (int i = 0; i < rowLen; i++) { if (buf[rowStart + i] != ' ' && buf[rowStart + i] != '\t') { isBlank = 0; break; } }
        if (isBlank || !memchr(buf + rowStart, '|', (size_t)rowLen)) break;

        int cellStart[TABLE_MAX_COLS], cellLen[TABLE_MAX_COLS];
        int rowCols = split_table_row(buf + rowStart, rowLen, cellStart, cellLen);
        consume_line(buf, len, pos);

        for (int c = 0; c < cols; c++) {
            /* Ragged rows (GFM spec): short rows are padded with empty cells, long rows are
               truncated -- tableCols is trusted unconditionally downstream (rendering,
               vertical-nav index arithmetic), so every row must end up with exactly `cols`
               cells. */
            const char *cellText = (c < rowCols) ? buf + rowStart + cellStart[c] : "";
            int cellTextLen = (c < rowCols) ? cellLen[c] : 0;
            document_insert_block(doc, doc->count, BLOCK_TABLE_CELL, cellText, cellTextLen);
            Block *b = &doc->blocks[doc->count - 1];
            b->tableCol = c;
            b->tableCols = cols;
            b->tableAlign = aligns[c];
        }
    }

    return 1;
}

/* ATX heading: 1-6 '#' followed by a space. Rejects 7+ '#' (not a heading per CommonMark). */
static int match_heading_prefix(const char *line, int len, int *level, int *prefixLen) {
    int n = 0;
    while (n < len && n < 7 && line[n] == '#') n++;
    if (n == 0 || n > 6) return 0;
    if (n >= len || line[n] != ' ') return 0;
    *level = n;
    *prefixLen = n + 1;
    return 1;
}

/* Task-list item: "- [ ] " / "- [x] " / "- [X] " (checked). Must be checked before the plain
   bullet prefix, which would otherwise match its first two bytes. */
static int match_task_prefix(const char *line, int len, int *prefixLen, int *checked) {
    if (len < 6) return 0;
    if ((line[0] != '-' && line[0] != '*') || line[1] != ' ') return 0;
    if (line[2] != '[' || line[4] != ']' || line[5] != ' ') return 0;
    char mark = line[3];
    if (mark != ' ' && mark != 'x' && mark != 'X') return 0;
    *checked = (mark != ' ');
    *prefixLen = 6;
    return 1;
}

bool document_load_file(Document *doc, const char *path) {
    FILE *f = fopen(path, "rb");
    if (!f) return false;

    fseek(f, 0, SEEK_END);
    long size = ftell(f);
    fseek(f, 0, SEEK_SET);
    if (size < 0) { fclose(f); return false; }

    char *buf = malloc((size_t)size + 1);
    size_t readN = fread(buf, 1, (size_t)size, f);
    fclose(f);
    buf[readN] = '\0';

    for (int i = 0; i < doc->count; i++) { sb_free(&doc->blocks[i].text); sb_free(&doc->blocks[i].lang); sb_free(&doc->blocks[i].alt); }
    doc->count = 0;

    int pos = 0;
    int len = (int)readN;
    while (pos < len) {
        int lineStart = pos;
        while (pos < len && buf[pos] != '\n') pos++;
        int lineEnd = pos;
        if (lineEnd > lineStart && buf[lineEnd - 1] == '\r') lineEnd--;
        if (pos < len) pos++; /* consume '\n' */

        const char *line = buf + lineStart;
        int lineLen = lineEnd - lineStart;

        if (is_code_fence(line, lineLen)) {
            int langStart = 3, langEnd = lineLen;
            while (langStart < langEnd && line[langStart] == ' ') langStart++;
            while (langEnd > langStart && line[langEnd - 1] == ' ') langEnd--;

            StrBuf code;
            sb_init(&code);
            int first = 1;
            while (pos < len) {
                int s = pos;
                while (pos < len && buf[pos] != '\n') pos++;
                int e = pos;
                if (e > s && buf[e - 1] == '\r') e--;
                if (pos < len) pos++;
                if (is_code_fence(buf + s, e - s)) break;
                if (!first) sb_append_char(&code, '\n');
                sb_append(&code, buf + s, e - s);
                first = 0;
            }
            document_insert_block(doc, doc->count, BLOCK_CODE, code.data, code.len);
            sb_free(&code);
            if (langEnd > langStart) sb_append(&doc->blocks[doc->count - 1].lang, line + langStart, langEnd - langStart);
            continue;
        }

        if (is_hr_line(line, lineLen)) {
            document_insert_block(doc, doc->count, BLOCK_HR, "", 0);
            continue;
        }

        int altStart, altLen, urlStart, urlLen;
        if (is_image_line(line, lineLen, &altStart, &altLen, &urlStart, &urlLen)) {
            document_insert_block(doc, doc->count, BLOCK_IMAGE, line + urlStart, urlLen);
            if (altLen > 0) sb_append(&doc->blocks[doc->count - 1].alt, line + altStart, altLen);
            continue;
        }

        if (is_blank_line(line, lineLen)) continue;

        int prefixLen, level, checked;
        if (match_heading_prefix(line, lineLen, &level, &prefixLen)) {
            static const BlockType headingTypes[6] = { BLOCK_H1, BLOCK_H2, BLOCK_H3, BLOCK_H4, BLOCK_H5, BLOCK_H6 };
            document_insert_block(doc, doc->count, headingTypes[level - 1], line + prefixLen, lineLen - prefixLen);
        } else if (match_task_prefix(line, lineLen, &prefixLen, &checked)) {
            document_insert_block(doc, doc->count, checked ? BLOCK_TASK_CHECKED : BLOCK_TASK_UNCHECKED, line + prefixLen, lineLen - prefixLen);
        } else if (lineLen >= 2 && (line[0] == '-' || line[0] == '*') && line[1] == ' ') {
            document_insert_block(doc, doc->count, BLOCK_BULLET, line + 2, lineLen - 2);
        } else if (lineLen >= 2 && line[0] == '>' && line[1] == ' ') {
            document_insert_block(doc, doc->count, BLOCK_QUOTE, line + 2, lineLen - 2);
        } else if (match_numbered_prefix(line, lineLen, &prefixLen)) {
            document_insert_block(doc, doc->count, BLOCK_NUMBERED, line + prefixLen, lineLen - prefixLen);
        } else if (memchr(line, '|', (size_t)lineLen) && try_parse_table(doc, buf, len, &pos, line, lineLen)) {
            /* try_parse_table already inserted the header+body cell blocks and advanced pos
               past every line it consumed. */
        } else {
            document_insert_block(doc, doc->count, BLOCK_PARAGRAPH, line, lineLen);
        }
    }

    free(buf);

    if (doc->count == 0) {
        document_insert_block(doc, 0, BLOCK_PARAGRAPH, "", 0);
    }

    return true;
}

/* Reconstructs the "| --- | :---: | ---: |"-style delimiter row for the table whose header row's
   last cell is at block index headerLastIdx, from each header cell's tableAlign. A plain 'l'
   (which also covers "no colons specified at all" -- both collapse to the same value on parse)
   serializes back as bare "---", not ":---" -- a narrow, accepted round-trip fidelity loss since
   both render identically. */
static void serialize_table_delimiter_row(StrBuf *out, Document *doc, int headerLastIdx, int cols) {
    int firstIdx = headerLastIdx - cols + 1;
    sb_append_char(out, '|');
    for (int c = 0; c < cols; c++) {
        char align = doc->blocks[firstIdx + c].tableAlign;
        sb_append(out, " ", 1);
        if (align == 'c') sb_append_char(out, ':');
        sb_append(out, "---", 3);
        if (align == 'c' || align == 'r') sb_append_char(out, ':');
        sb_append(out, " |", 2);
    }
}

void document_serialize(Document *doc, StrBuf *out) {
    sb_clear(out);

    int numberedRun = 0;
    for (int i = 0; i < doc->count; i++) {
        Block *b = &doc->blocks[i];

        int skipDefaultSeparator = 0;
        if (i > 0) {
            Block *prev = &doc->blocks[i - 1];
            if (block_type_is_table_cell(prev->type) && block_type_is_table_cell(b->type) && prev->tableCols == b->tableCols) {
                if (b->tableCol != 0) {
                    sb_append(out, " | ", 3);
                } else {
                    sb_append_char(out, '\n');
                    if (prev->type == BLOCK_TABLE_HEADER_CELL) {
                        serialize_table_delimiter_row(out, doc, i - 1, prev->tableCols);
                        sb_append_char(out, '\n');
                    }
                }
                skipDefaultSeparator = 1;
            }
        }
        if (i > 0 && !skipDefaultSeparator) sb_append(out, "\n\n", 2);

        if (block_type_is_table_cell(b->type) && b->tableCol == 0) sb_append(out, "| ", 2);

        switch (b->type) {
            case BLOCK_H1: sb_append(out, "# ", 2); sb_append(out, b->text.data, b->text.len); break;
            case BLOCK_H2: sb_append(out, "## ", 3); sb_append(out, b->text.data, b->text.len); break;
            case BLOCK_H3: sb_append(out, "### ", 4); sb_append(out, b->text.data, b->text.len); break;
            case BLOCK_H4: sb_append(out, "#### ", 5); sb_append(out, b->text.data, b->text.len); break;
            case BLOCK_H5: sb_append(out, "##### ", 6); sb_append(out, b->text.data, b->text.len); break;
            case BLOCK_H6: sb_append(out, "###### ", 7); sb_append(out, b->text.data, b->text.len); break;
            case BLOCK_BULLET: sb_append(out, "- ", 2); sb_append(out, b->text.data, b->text.len); break;
            case BLOCK_TASK_UNCHECKED: sb_append(out, "- [ ] ", 6); sb_append(out, b->text.data, b->text.len); break;
            case BLOCK_TASK_CHECKED: sb_append(out, "- [x] ", 6); sb_append(out, b->text.data, b->text.len); break;
            case BLOCK_NUMBERED: {
                if (i > 0 && doc->blocks[i - 1].type == BLOCK_NUMBERED) numberedRun++;
                else numberedRun = 1;
                char prefix[16];
                int n = snprintf(prefix, sizeof prefix, "%d. ", numberedRun);
                sb_append(out, prefix, n);
                sb_append(out, b->text.data, b->text.len);
                break;
            }
            case BLOCK_QUOTE: sb_append(out, "> ", 2); sb_append(out, b->text.data, b->text.len); break;
            case BLOCK_HR: sb_append(out, "---", 3); break;
            case BLOCK_IMAGE:
                sb_append(out, "![", 2);
                sb_append(out, b->alt.data, b->alt.len);
                sb_append(out, "](", 2);
                sb_append(out, b->text.data, b->text.len);
                sb_append_char(out, ')');
                break;
            case BLOCK_CODE:
                sb_append(out, "```", 3);
                sb_append(out, b->lang.data, b->lang.len);
                sb_append_char(out, '\n');
                sb_append(out, b->text.data, b->text.len);
                sb_append(out, "\n```", 4);
                break;
            case BLOCK_TABLE_HEADER_CELL:
            case BLOCK_TABLE_CELL:
                sb_append(out, b->text.data, b->text.len);
                break;
            case BLOCK_PARAGRAPH:
            default:
                sb_append(out, b->text.data, b->text.len);
                break;
        }

        if (block_type_is_table_cell(b->type) && b->tableCol == b->tableCols - 1) sb_append(out, " |", 2);
    }
    sb_append_char(out, '\n');
}

bool document_save_file(Document *doc, const char *path) {
    StrBuf out;
    sb_init(&out);
    document_serialize(doc, &out);

    FILE *f = fopen(path, "wb");
    if (!f) { sb_free(&out); return false; }
    size_t written = fwrite(out.data, 1, (size_t)out.len, f);
    fclose(f);
    sb_free(&out);
    return written == (size_t)out.len || out.len == 0;
}

void document_compute_stats(Document *doc, DocStats *out) {
    StrBuf tmp;
    sb_init(&tmp);
    document_serialize(doc, &tmp);
    out->byteSize = tmp.len;
    int lines = 0;
    for (int i = 0; i < tmp.len; i++) if (tmp.data[i] == '\n') lines++;
    out->lineCount = lines;
    sb_free(&tmp);
}