#include "tokenizer.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <ctype.h>

static unsigned int hash_str(const char* s) {
    unsigned int h = 5381;
    while (*s) h = ((h << 5) + h) + (unsigned char)(*s++);
    return h;
}

static void hash_insert(Tokenizer* tok, const char* token, int id) {
    unsigned int h = hash_str(token) % (unsigned int)tok->hash_size;
    while (tok->hash_keys[h] != -1)
        h = (h + 1) % (unsigned int)tok->hash_size;
    tok->hash_keys[h] = 1;
    tok->hash_vals[h] = id;
    strncpy(tok->vocab[id].token, token, TOK_MAX_WORD_LEN - 1);
}

static int hash_lookup(const Tokenizer* tok, const char* token) {
    unsigned int h = hash_str(token) % (unsigned int)tok->hash_size;
    int probes = 0;
    while (tok->hash_keys[h] != -1 && probes < tok->hash_size) {
        int id = tok->hash_vals[h];
        if (strncmp(tok->vocab[id].token, token, TOK_MAX_WORD_LEN) == 0)
            return id;
        h = (h + 1) % (unsigned int)tok->hash_size;
        probes++;
    }
    return TOK_UNK;
}

static void add_special_tokens(Tokenizer* tok) {
    const char* specials[] = {"<unk>","<bos>","<eos>","<pad>","<sep>"};
    for (int i = 0; i < 5; i++) {
        strncpy(tok->vocab[i].token, specials[i], TOK_MAX_WORD_LEN - 1);
        tok->vocab[i].freq = 999999;
        hash_insert(tok, specials[i], i);
    }
    tok->vocab_size = 5;
}

static int extract_field(const char* line, const char* key, char* out, int out_len) {
    char search[64];
    snprintf(search, sizeof(search), "\"%s\"", key);
    const char* p = strstr(line, search);
    if (!p) return 0;
    p += strlen(search);
    while (*p && (*p == ':' || *p == ' ')) p++;
    if (*p != '"') return 0;
    p++;
    int i = 0;
    while (*p && *p != '"' && i < out_len - 1) {
        if (*p == '\\') { p++; if (*p) { out[i++] = *p++; } }
        else out[i++] = *p++;
    }
    out[i] = '\0';
    return i > 0;
}

static void count_tokens(Tokenizer* tok, const char* text) {
    char word[TOK_MAX_WORD_LEN];
    int wi = 0;
    const char* p = text;
    while (*p) {
        if (isalnum((unsigned char)*p) || *p == '\'') {
            if (wi < TOK_MAX_WORD_LEN - 1)
                word[wi++] = tolower((unsigned char)*p);
            p++;
        } else {
            if (wi > 0) {
                word[wi] = '\0';
                int found = 0;
                for (int i = 5; i < tok->vocab_size; i++) {
                    if (strncmp(tok->vocab[i].token, word, TOK_MAX_WORD_LEN) == 0) {
                        tok->vocab[i].freq++; found = 1; break;
                    }
                }
                if (!found && tok->vocab_size < TOK_VOCAB_SIZE) {
                    strncpy(tok->vocab[tok->vocab_size].token, word, TOK_MAX_WORD_LEN - 1);
                    tok->vocab[tok->vocab_size].freq = 1;
                    tok->vocab_size++;
                }
                wi = 0;
            }
            if (!isspace((unsigned char)*p)) {
                char punct[3] = {*p, '\0', '\0'};
                int found = 0;
                for (int i = 5; i < tok->vocab_size; i++) {
                    if (strncmp(tok->vocab[i].token, punct, TOK_MAX_WORD_LEN) == 0) {
                        tok->vocab[i].freq++; found = 1; break;
                    }
                }
                if (!found && tok->vocab_size < TOK_VOCAB_SIZE) {
                    strncpy(tok->vocab[tok->vocab_size].token, punct, TOK_MAX_WORD_LEN - 1);
                    tok->vocab[tok->vocab_size].freq = 1;
                    tok->vocab_size++;
                }
            }
            p++;
        }
    }
    if (wi > 0) {
        word[wi] = '\0';
        int found = 0;
        for (int i = 5; i < tok->vocab_size; i++) {
            if (strncmp(tok->vocab[i].token, word, TOK_MAX_WORD_LEN) == 0) {
                tok->vocab[i].freq++; found = 1; break;
            }
        }
        if (!found && tok->vocab_size < TOK_VOCAB_SIZE) {
            strncpy(tok->vocab[tok->vocab_size].token, word, TOK_MAX_WORD_LEN - 1);
            tok->vocab[tok->vocab_size].freq = 1;
            tok->vocab_size++;
        }
    }
}

static int cmp_freq(const void* a, const void* b) {
    return ((TokenEntry*)b)->freq - ((TokenEntry*)a)->freq;
}

int tokenizer_build(Tokenizer* tok, const char* jsonl_path, int max_lines) {
    tok->hash_size = TOK_VOCAB_SIZE * 4;
    memset(tok->hash_keys, -1, sizeof(tok->hash_keys));
    memset(tok->hash_vals, -1, sizeof(tok->hash_vals));
    tok->vocab_size = 0;
    add_special_tokens(tok);
    FILE* f = fopen(jsonl_path, "r");
    if (!f) { fprintf(stderr, "[TOK] cannot open %s\n", jsonl_path); return 0; }
    char line[4096], field[2048];
    int lines = 0;
    printf("[TOK] Building vocab from %s ...\n", jsonl_path);
    while (fgets(line, sizeof(line), f) && (max_lines <= 0 || lines < max_lines)) {
        if (extract_field(line, "input",  field, sizeof(field))) count_tokens(tok, field);
        if (extract_field(line, "output", field, sizeof(field))) count_tokens(tok, field);
        lines++;
        if (lines % 50000 == 0)
            printf("[TOK] Processed %d lines, vocab=%d\n", lines, tok->vocab_size);
        
    }
    fclose(f);
    qsort(tok->vocab + 5, tok->vocab_size - 5, sizeof(TokenEntry), cmp_freq);
    if (tok->vocab_size > TOK_VOCAB_SIZE) tok->vocab_size = TOK_VOCAB_SIZE;
    memset(tok->hash_keys, -1, sizeof(tok->hash_keys));
    for (int i = 0; i < tok->vocab_size; i++)
        hash_insert(tok, tok->vocab[i].token, i);
    printf("[TOK] Final vocab size: %d (from %d lines)\n", tok->vocab_size, lines);
    return tok->vocab_size;
}

int tokenizer_save(const Tokenizer* tok, const char* path) {
    FILE* f = fopen(path, "w");
    if (!f) return 0;
    fprintf(f, "%d\n", tok->vocab_size);
    for (int i = 0; i < tok->vocab_size; i++)
        fprintf(f, "%s\t%d\n", tok->vocab[i].token, tok->vocab[i].freq);
    fclose(f);
    printf("[TOK] Saved vocab to %s\n", path);
    return 1;
}

int tokenizer_load(Tokenizer* tok, const char* path) {
    FILE* f = fopen(path, "r");
    if (!f) { fprintf(stderr, "[TOK] cannot open %s\n", path); return 0; }
    tok->hash_size = TOK_VOCAB_SIZE * 4;
    memset(tok->hash_keys, -1, sizeof(tok->hash_keys));
    memset(tok->hash_vals, -1, sizeof(tok->hash_vals));
    fscanf(f, "%d\n", &tok->vocab_size);
    for (int i = 0; i < tok->vocab_size; i++) {
        fscanf(f, "%31s\t%d\n", tok->vocab[i].token, &tok->vocab[i].freq);
        hash_insert(tok, tok->vocab[i].token, i);
    }
    fclose(f);
    printf("[TOK] Loaded vocab: %d tokens from %s\n", tok->vocab_size, path);
    return 1;
}

int tokenizer_encode(const Tokenizer* tok, const char* text, int* buf, int buf_len) {
    int n = 0;
    char word[TOK_MAX_WORD_LEN];
    int wi = 0;
    const char* p = text;
    while (*p && n < buf_len) {
        if (isalnum((unsigned char)*p) || *p == '\'') {
            if (wi < TOK_MAX_WORD_LEN - 1)
                word[wi++] = tolower((unsigned char)*p);
            p++;
        } else {
            if (wi > 0) {
                word[wi] = '\0';
                buf[n++] = hash_lookup(tok, word);
                wi = 0;
            }
            if (!isspace((unsigned char)*p) && n < buf_len) {
                char punct[3] = {*p, '\0', '\0'};
                buf[n++] = hash_lookup(tok, punct);
            }
            p++;
        }
    }
    if (wi > 0 && n < buf_len) {
        word[wi] = '\0';
        buf[n++] = hash_lookup(tok, word);
    }
    return n;
}

void tokenizer_decode(const Tokenizer* tok, const int* ids, int n, char* out, int out_len) {
    int pos = 0;
    for (int i = 0; i < n && pos < out_len - 2; i++) {
        if (ids[i] == TOK_EOS || ids[i] == TOK_PAD) break;
        if (ids[i] == TOK_BOS || ids[i] == TOK_SEP) continue;
        const char* tok_str = tokenizer_id_to_token(tok, ids[i]);
        int len = strlen(tok_str);
        if (i > 0 && isalpha((unsigned char)tok_str[0]) && pos < out_len - 2)
            out[pos++] = ' ';
        if (pos + len < out_len - 1) {
            memcpy(out + pos, tok_str, len);
            pos += len;
        }
    }
    out[pos] = '\0';
}

int tokenizer_token_to_id(const Tokenizer* tok, const char* token) {
    return hash_lookup(tok, token);
}

const char* tokenizer_id_to_token(const Tokenizer* tok, int id) {
    if (id < 0 || id >= tok->vocab_size) return "<unk>";
    return tok->vocab[id].token;
}