#include "common.h"
#include <assert.h>
#include <errno.h>
#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

typedef struct {
    char *text;
    uint32_t len;
    uint32_t count;
} MultisetEntry;

typedef struct {
    MultisetEntry *entries;
    size_t entries_cap;
    size_t entries_size;
} Multiset;

static int multisetentry_cmp(const MultisetEntry *a, const MultisetEntry *b)
{
    if (a->len != b->len)
        return (int)a->len - (int)b->len;
    return strcmp(a->text, b->text);
}

static int multisetentry_cmp_wordcount(
    const MultisetEntry *a, const MultisetEntry *b)
{
    return (int)a->count - (int)b->count;
}

static void multimap_search(Multiset *s,
    const MultisetEntry *entry,
    MultisetEntry **out_entry,
    size_t *out_idx)
{
    *out_entry = binary_search(s->entries,
        s->entries_size,
        (void *)entry,
        sizeof(*entry),
        out_idx,
        (CmpFn)multisetentry_cmp);
}

void multiset_free(Multiset *s)
{
    for (size_t i = 0; i < s->entries_size; ++i)
        free(s->entries[i].text);
    if (s->entries)
        free(s->entries);
}

void multiset_dump(Multiset *s)
{
    printf("entries_size = %lu\n", s->entries_size);
    for (size_t i = 0; i < s->entries_size; ++i)
        printf("  %s -> %u\n", s->entries[i].text, s->entries[i].count);
}

// strdup was introduced in C23, we're using C17
#if (__STDC_VERSION__ < 202311L)
char *strdup(const char *str)
{
    char *p = malloc(strlen(str) + 1);
    strcpy(p, str);
    return p;
}
#endif

void multiset_insert(Multiset *s, const char *text)
{
    if (!s->entries_size) {
        MultisetEntry entry = {
            .text = strdup(text),
            .len = strlen(text),
            .count = 1,
        };
        array_push((void **)&s->entries,
            &s->entries_cap,
            &s->entries_size,
            &entry,
            sizeof(entry));
        return;
    }

    MultisetEntry search_entry = {
        .text = (char *)text, // NOTE: cast away const. we'll borrow it shortly
        .len = strlen(text),
        .count = 0,
    };

    MultisetEntry *out_entry;
    size_t out_idx;
    multimap_search(s, &search_entry, &out_entry, &out_idx);

    if (out_entry) {
        out_entry->count += 1;
        return;
    }

    MultisetEntry last_entry = s->entries[s->entries_size - 1];
    array_push((void **)&s->entries,
        &s->entries_cap,
        &s->entries_size,
        &last_entry,
        sizeof(last_entry));

    memmove(&s->entries[out_idx + 1],
        &s->entries[out_idx],
        sizeof(s->entries[0]) * (s->entries_size - out_idx - 1));

    s->entries[out_idx] = (MultisetEntry) {
        .text = strdup(text),
        .len = strlen(text),
        .count = 1,
    };
}

void test_multiset(void);

void multiset_sort(Multiset *s)
{
    qsort(s->entries,
        s->entries_size,
        sizeof(s->entries[0]),
        (CmpFn)multisetentry_cmp);
}

void multiset_sort_wordcount(Multiset *s)
{
    qsort(s->entries,
        s->entries_size,
        sizeof(s->entries[0]),
        (CmpFn)multisetentry_cmp_wordcount);
}

static int char_cmp(const char *a, const char *b)
{
    return *a - *b;
}

typedef struct {
    uint16_t left, right;
} Tok;

static int tok_cmp(const Tok *a, const Tok *b)
{
    return a->left != b->left ? a->left - b->left : a->right - b->right;
}

typedef struct {
    const char *chars;
    size_t chars_size;
    Tok *toks;
    size_t toks_cap;
    size_t toks_size;
} Vocab;

uint16_t vocab_add_tok(Vocab *v, Tok tok)
{
    uint16_t tokval = (uint16_t)v->toks_size + v->chars_size + 1;

    array_push(
        (void **)&v->toks, &v->toks_cap, &v->toks_size, &tok, sizeof(tok));

    return tokval;
}

uint16_t vocab_char_tokval(Vocab *v, char ch)
{
    size_t idx;
    char *found = binary_search((void *)v->chars,
        v->chars_size,
        &ch,
        sizeof(ch),
        &idx,
        (CmpFn)char_cmp);

    assert(found);
    return (uint16_t)idx + 1;
}

size_t vocab_tok_str(
    Vocab *v, char str[], size_t str_size, const Tok *tok, bool subst_space);

size_t vocab_tokval_str(
    Vocab *v, char str[], size_t str_size, uint16_t tokval, bool subst_space)
{
    if (tokval == 0) {
        assert(str_size >= 1);
        str[0] = '0';
        return 1;
    }

    if (tokval < v->chars_size + 1) {
        char ch = v->chars[tokval - 1];

        if (ch == ' ') {
            if (subst_space) {
                char *blank = "␣";
                size_t blank_len = strlen(blank);
                assert(str_size >= blank_len);
                memcpy(str, blank, blank_len);
                return blank_len;
            }

            assert(str_size >= 1);
            str[0] = ' ';
            return 1;
        }

        assert(str_size >= 1);
        str[0] = v->chars[tokval - 1];
        return 1;
    }

    return vocab_tok_str(
        v, str, str_size, &v->toks[tokval - v->chars_size - 1], subst_space);
}

size_t vocab_tok_str(
    Vocab *v, char str[], size_t str_size, const Tok *tok, bool subst_space)
{
    size_t len = 0;
    len += vocab_tokval_str(v, str, str_size, tok->left, subst_space);
    len += vocab_tokval_str(
        v, str + len, str_size - len, tok->right, subst_space);
    return len;
}

void vocab_print_tokval(Vocab *v, FILE *file, uint16_t tokval, bool subst_space)
{
    char tokstr[60];
    size_t tokstr_len
        = vocab_tokval_str(v, tokstr, sizeof(tokstr) - 1, tokval, subst_space);
    tokstr[tokstr_len] = '\0';
    fprintf(file, "%s", tokstr);
}

typedef struct {
    size_t size_of_word;
    uint16_t *words;
    size_t words_size;
} Corpus;

Tok corpus_most_common_tok_pair(Corpus *c, Vocab *v)
{
    (void)v;

    typedef struct {
        Tok tok;
        uint32_t count;
    } TokPair;

    TokPair *pairs = NULL;
    size_t pairs_cap = 0;
    size_t pairs_size = 0;

    for (size_t word_idx = 0; word_idx < c->words_size; ++word_idx) {
        uint16_t *word = &c->words[c->size_of_word * word_idx];

        for (size_t i = 0; i < c->size_of_word - 1 && word[i + 1]; ++i) {
            Tok tok = { word[i], word[i + 1] };

            if (!pairs_size) {
                TokPair pair = { tok, 1 };
                array_push((void **)&pairs,
                    &pairs_cap,
                    &pairs_size,
                    &pair,
                    sizeof(pair));
                continue;
            }

            size_t idx;
            TokPair *found = binary_search(pairs,
                pairs_size,
                &tok,
                sizeof(pairs[0]),
                &idx,
                (CmpFn)tok_cmp);

            if (found) {
                found->count += 1;
                continue;
            }

            TokPair empty_pair = { 0 };
            array_push((void **)&pairs,
                &pairs_cap,
                &pairs_size,
                &empty_pair,
                sizeof(empty_pair));

            memmove(&pairs[idx + 1],
                &pairs[idx],
                sizeof(pairs[0]) * (pairs_size - idx - 1));

            pairs[idx] = (TokPair) { tok, 1 };
        }
    }

    TokPair *max_pair = &pairs[0];
    for (size_t i = 0; i < pairs_size; ++i) {
        if (pairs[i].count > max_pair->count)
            max_pair = &pairs[i];
    }

    Tok max_tok = max_pair->tok;
    free(pairs);
    return max_tok;
}

void corpus_rewrite_with_tok(Corpus *c, Tok tok, Vocab *v)
{
    uint16_t tok_idx = vocab_add_tok(v, tok);

    for (size_t word_idx = 0; word_idx < c->words_size; ++word_idx) {
        uint16_t *word = &c->words[c->size_of_word * word_idx];

        for (size_t i = 0; i < c->size_of_word - 1 && word[i + 1]; ++i) {
            if (word[i] != tok.left || word[i + 1] != tok.right)
                continue;

            memmove(&word[i + 1],
                &word[i + 2],
                sizeof(uint16_t) * (c->size_of_word - i - 2));

            word[i] = tok_idx;

            // NOTE: since the token is newly introduced,
            // we don't need to check it again,
            // hence we don't set i back one here.
        }
    }
}

static void corpus_dump_word(Corpus *c, Vocab *v, size_t word_idx)
{
    uint16_t *word = &c->words[c->size_of_word * word_idx];

    printf("[ ");
    for (size_t i = 0; i < c->size_of_word; ++i) {
        if (!word[i])
            break;

        vocab_print_tokval(v, stdout, word[i], true);
        printf(" ");
    }
    printf("]");
}

void corpus_dump(Corpus *c, Vocab *v)
{
    size_t batch_size = 10;
    for (size_t word_idx = 0;
        word_idx < c->words_size && word_idx < batch_size * 10;
        ++word_idx) {
        corpus_dump_word(c, v, word_idx);
        printf("\n");
    }
    printf("...\n");
    for (size_t word_idx = c->words_size / 2 - batch_size / 2;
        word_idx < c->words_size
        && word_idx < c->words_size / 2 + batch_size / 2;
        ++word_idx) {
        corpus_dump_word(c, v, word_idx);
        printf("\n");
    }
    printf("...\n");
    // for (size_t word_idx = c->words_size - batch_size; word_idx <
    // c->words_size;
    //     ++word_idx) {
    //     corpus_dump_word(c, v, word_idx);
    //     printf("\n");
    // }
}

int main(int argc, char *argv[])
{
    test_multiset();

    if (argc <= 2) {
        fprintf(stderr, "error: no files specified\n");
        fprintf(
            stderr, "usage: tokenizer <input text file> <output token file>\n");
        return 1;
    }

    const char *input_filename = argv[1];
    const char *output_filename = argv[2];

    FILE *input_file = fopen(input_filename, "r");
    if (!input_file) {
        fprintf(stderr,
            "error: could not open file '%s': %s\n",
            input_filename,
            strerror(errno));
        return 1;
    }

    Multiset words = { NULL, 0, 0 };
    size_t total_word_count = 0;
    size_t max_word_len = 0;

    char line[64];
    size_t line_size = sizeof(line) / sizeof(line[0]);
    size_t len;

    printf("reading '%s'...\n", input_filename);

    int ch = 0;
    while ((len = next_word(line, line_size, input_file, &ch))) {
        line[len] = '\0';
        // printf("%s\n", line);

        if (line[0] == '\n') {
            assert(len == 1);
            continue;
        }

        multiset_insert(&words, line);
        total_word_count += 1;

        if (len > max_word_len)
            max_word_len = len;
    }

    fclose(input_file);

    printf("total word count: %ld\n", total_word_count);
    printf("unique word count: %ld\n", words.entries_size);
    printf("max word length: %ld\n", max_word_len);

    multiset_sort_wordcount(&words);

    size_t selected_words_count = 50000;

    struct {
        char *data;
        size_t cap;
        size_t size;
    } vocab_chars = { NULL, 0, 0 };

    printf("discovering initial vocabulary...\n");

    for (size_t i = 0; i < words.entries_size && i < selected_words_count;
        ++i) {
        MultisetEntry *entry = &words.entries[words.entries_size - 1 - i];

        for (size_t j = 0; j < entry->len; ++j) {
            char ch = entry->text[j];

            size_t idx = 0;
            char *found = binary_search(vocab_chars.data,
                vocab_chars.size,
                &ch,
                sizeof(ch),
                &idx,
                (CmpFn)char_cmp);

            if (found)
                continue;

            array_push((void **)&vocab_chars.data,
                &vocab_chars.cap,
                &vocab_chars.size,
                &ch,
                sizeof(ch));

            if (vocab_chars.size == 1)
                continue;

            memmove(&vocab_chars.data[idx + 1],
                &vocab_chars.data[idx],
                sizeof(vocab_chars.data[0]) * (vocab_chars.size - idx - 1));

            vocab_chars.data[idx] = ch;
        }
    }

    printf("[");
    for (size_t i = 0; i < vocab_chars.size; ++i)
        printf("%c", vocab_chars.data[i]);
    printf("]\n");

    Vocab vocab = {
        .chars = vocab_chars.data,
        .chars_size = vocab_chars.size,
        .toks = NULL,
        .toks_cap = 0,
        .toks_size = 0,
    };

    size_t word_size_aligned = (max_word_len + 3) & ~3ul;

    size_t corpus_words_cap = selected_words_count;
    uint16_t *corpus_words
        = calloc(corpus_words_cap, word_size_aligned * sizeof(uint16_t));

    Corpus corpus = {
        .size_of_word = word_size_aligned,
        .words = corpus_words,
        .words_size = 0,
    };

    printf("using word size = %ld, word count = %ld\n",
        corpus.size_of_word,
        corpus.size_of_word);
    printf("converting words to tokens...\n");

    for (size_t word_idx = 0; word_idx < selected_words_count; ++word_idx) {

        assert(word_idx < words.entries_size);

        MultisetEntry entry = words.entries[words.entries_size - word_idx - 1];
        assert(entry.len != 0 && entry.text[0]);

        for (size_t i = 0; i < corpus.size_of_word && i < entry.len; ++i) {
            corpus.words[word_idx * word_size_aligned + i]
                = vocab_char_tokval(&vocab, entry.text[i]);
        }

        corpus.words_size += 1;
    }

    multiset_free(&words);

    size_t comb_iters = 4000;

    printf("combining tokens in %ld iterations...\n", comb_iters);

    for (size_t j = 0; j < comb_iters / 100; ++j) {
        for (size_t i = 0; i < 100; ++i) {
            Tok tok = corpus_most_common_tok_pair(&corpus, &vocab);
            corpus_rewrite_with_tok(&corpus, tok, &vocab);
        }

        printf("%ld / %ld, vocabulary size = %ld\n",
            (j + 1) * 100,
            comb_iters,
            vocab.chars_size + vocab.toks_size + 1);
    }

    corpus_dump(&corpus, &vocab);

    FILE *output_file = fopen(output_filename, "w");
    if (!output_file) {
        fprintf(stderr,
            "error: could not open file '%s' for writing: %s\n",
            output_filename,
            strerror(errno));
        return 1;
    }

    for (size_t i = 0; i < vocab_chars.size; ++i) {
        fputc(vocab_chars.data[i], output_file);
    }
    fputc('\n', output_file);

    for (size_t i = 0; i < vocab.toks_size; ++i) {
        vocab_print_tokval(
            &vocab, output_file, 1 + vocab.chars_size + i, false);
        fputc('\n', output_file);
    }

    fclose(output_file);

    free(vocab_chars.data);
    free(vocab.toks);
    free(corpus.words);
}

void test_multiset(void)
{
    {
        Multiset m = { NULL, 0, 0 };

        multiset_insert(&m, "a");
        assert(m.entries_size == 1);
        assert(strcmp(m.entries[0].text, "a") == 0);
        assert(m.entries[0].count == 1);

        multiset_insert(&m, "d");
        assert(m.entries_size == 2);
        assert(strcmp(m.entries[0].text, "a") == 0);
        assert(m.entries[0].count == 1);
        assert(strcmp(m.entries[1].text, "d") == 0);
        assert(m.entries[1].count == 1);

        multiset_insert(&m, "d");
        assert(m.entries_size == 2);
        assert(m.entries[1].count == 2);

        multiset_insert(&m, "b");
        assert(m.entries_size == 3);
        assert(strcmp(m.entries[0].text, "a") == 0);
        assert(strcmp(m.entries[1].text, "b") == 0);
        assert(strcmp(m.entries[2].text, "d") == 0);

        multiset_insert(&m, ".");
        assert(m.entries_size == 4);
        assert(strcmp(m.entries[0].text, ".") == 0);

        multiset_insert(&m, "aa");
        assert(m.entries_size == 5);
        assert(strcmp(m.entries[4].text, "aa") == 0);

        multiset_insert(&m, "c");
        multiset_insert(&m, "c");
        multiset_insert(&m, "c");
        assert(m.entries_size == 6);
        assert(strcmp(m.entries[3].text, "c") == 0);
        assert(m.entries[3].count == 3);

        multiset_free(&m);
    }
}