#include "loader.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

/* extract value of a JSON string field */
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 == 'n')       { out[i++] = ' '; p++; }
            else if (*p == 't')  { out[i++] = ' '; p++; }
            else if (*p == '"')  { out[i++] = '"'; p++; }
            else if (*p == '\\') { out[i++] = '\\'; p++; }
            else                 { out[i++] = *p++; }
        } else {
            out[i++] = *p++;
        }
    }
    out[i] = '\0';
    return i > 0;
}

int loader_open(DataLoader* dl, const char* path, Tokenizer* tok) {
    dl->file = fopen(path, "r");
    if (!dl->file) {
        fprintf(stderr, "[LOADER] Cannot open %s\n", path);
        return 0;
    }
    dl->path         = (char*)path;
    dl->tok          = tok;
    dl->current_line = 0;
    dl->total_lines  = 0;
    printf("[LOADER] Opened %s\n", path);
    return 1;
}

void loader_close(DataLoader* dl) {
    if (dl->file) { fclose(dl->file); dl->file = NULL; }
}

void loader_rewind(DataLoader* dl) {
    if (dl->file) rewind(dl->file);
    dl->current_line = 0;
}

long loader_count_lines(const char* path) {
    FILE* f = fopen(path, "r");
    if (!f) return 0;
    long count = 0;
    int c;
    while ((c = fgetc(f)) != EOF) if (c == '\n') count++;
    fclose(f);
    return count;
}

void loader_skip(DataLoader* dl, long n) {
    char line[8192];
    for (long i = 0; i < n; i++) {
        if (!fgets(line, sizeof(line), dl->file)) break;
        dl->current_line++;
    }
}

int loader_next(DataLoader* dl, int* token_buf, int buf_len) {
    if (!dl->file) return 0;

    char line[8192];
    char input_text[2048];
    char output_text[2048];

    /* skip lines that don't have both input and output */
    while (fgets(line, sizeof(line), dl->file)) {
        dl->current_line++;

        int has_input  = extract_field(line, "input",  input_text,  sizeof(input_text));
        int has_output = extract_field(line, "output", output_text, sizeof(output_text));

        if (!has_input || !has_output) continue;
        if (strlen(input_text) < 2 || strlen(output_text) < 2) continue;

        /* build token sequence:
           [BOS] input_tokens [SEP] output_tokens [EOS] */
        int n = 0;

        token_buf[n++] = TOK_BOS;

        /* encode input */
        int input_ids[JOY_SEQLEN];
        int ni = tokenizer_encode(dl->tok, input_text, input_ids, JOY_SEQLEN / 2);
        for (int i = 0; i < ni && n < buf_len - 3; i++)
            token_buf[n++] = input_ids[i];

        token_buf[n++] = TOK_SEP;

        /* encode output */
        int output_ids[JOY_SEQLEN];
        int no = tokenizer_encode(dl->tok, output_text, output_ids, JOY_SEQLEN / 2);
        for (int i = 0; i < no && n < buf_len - 1; i++)
            token_buf[n++] = output_ids[i];

        token_buf[n++] = TOK_EOS;

        return n;
    }

    return 0; /* end of file */
}