#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <math.h>
#include <string.h>
#include <time.h>
#include <vector>
#include <algorithm>
#include <immintrin.h> 
#include <chrono>
#include <assert.h>

#ifdef _WIN32
#include <windows.h>
#include <psapi.h>
#else
#include <sys/mman.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#endif

#ifdef _OPENMP
#include <omp.h>
#endif

#pragma pack(push, 1)
typedef struct {
    char magic[4];        // "NANO"
    uint32_t version;     // 1 (untied) or 2 (tied)
    uint32_t header_size; // 64
    int n_layer;
    int n_head;
    int n_embd;
    int block_size;
    int vocab_size;
    char padding[32];
} ModelHeader;
#pragma pack(pop)

typedef struct {
    int n_layer;
    int n_head;
    int n_embd;
    int block_size;
    int vocab_size;
} Config;

typedef struct {
    float* wte; float* wpe;
    float** ln1_w; float** ln1_b;
    float** c_attn_w; float** c_attn_b;
    float** c_proj_w; float** c_proj_b;
    float** ln2_w; float** ln2_b;
    float** fc_w; float** fc_b;
    float** mlp_proj_w; float** mlp_proj_b;
    float* ln_f_w; float* ln_f_b;
    float* lm_head_w;
} Weights;

typedef struct { float* k_cache; float* v_cache; } KVCache;

static Config cfg;
static Weights W;
static float* model_data_buffer = NULL;

// Operation profiling accumulators (microseconds)
static long long prof_attn_us = 0;
static long long prof_mlp_us = 0;
static long long prof_lmhead_us = 0;
static long long prof_sample_us = 0;

void print_mem_stage(const char* stage) {
#ifdef _WIN32
    PROCESS_MEMORY_COUNTERS_EX pmc;
    if (GetProcessMemoryInfo(GetCurrentProcess(), (PROCESS_MEMORY_COUNTERS*)&pmc, sizeof(pmc))) {
        double ws_mb = (double)pmc.WorkingSetSize / (1024.0 * 1024.0);
        double priv_mb = (double)pmc.PrivateUsage / (1024.0 * 1024.0);
        double shared_ws_mb = ws_mb > priv_mb ? (ws_mb - priv_mb) : 0.0;
        fprintf(stderr, "STAGE_MEM: stage=%-28s ws_mb=%7.2f priv_mb=%7.2f shared_mb=%7.2f\n",
                stage, ws_mb, priv_mb, shared_ws_mb);
        fflush(stderr);
    }
#else
    FILE* f = fopen("/proc/self/statm", "r");
    if (f) {
        long total_pages = 0, resident_pages = 0, shared_pages = 0;
        if (fscanf(f, "%ld %ld %ld", &total_pages, &resident_pages, &shared_pages) == 3) {
            long page_size = sysconf(_SC_PAGESIZE);
            double ws_mb = (double)(resident_pages * page_size) / (1024.0 * 1024.0);
            double shared_mb = (double)(shared_pages * page_size) / (1024.0 * 1024.0);
            double priv_mb = ws_mb > shared_mb ? (ws_mb - shared_mb) : 0.0;
            fprintf(stderr, "STAGE_MEM: stage=%-28s ws_mb=%7.2f priv_mb=%7.2f shared_mb=%7.2f\n",
                    stage, ws_mb, priv_mb, shared_mb);
            fflush(stderr);
        }
        fclose(f);
    }
#endif
}

static void layer_norm(float* out, const float* x, const float* w, const float* b, int size) {
    float mean = 0.0f, var = 0.0f;
    for (int i = 0; i < size; i++) mean += x[i];
    mean /= size;
    for (int i = 0; i < size; i++) {
        float d = x[i] - mean;
        var += d * d;
    }
    var /= size;
    float scale = 1.0f / sqrtf(var + 1e-5f);
    for (int i = 0; i < size; i++)
        out[i] = (x[i] - mean) * scale * w[i] + b[i];
}

static void matmul_vec(float* out, const float* mat, const float* x, int M, int K) {
#pragma omp parallel for
    for (int i = 0; i < M; i++) {
        const float* row = mat + (long long)i * K;
        __m256 sum_vec = _mm256_setzero_ps();
        int j = 0;
        for (; j <= K - 8; j += 8) {
            __m256 m_val = _mm256_loadu_ps(&row[j]);
            __m256 x_val = _mm256_loadu_ps(&x[j]);
            sum_vec = _mm256_fmadd_ps(m_val, x_val, sum_vec);
        }
        float sum_arr[8];
        _mm256_storeu_ps(sum_arr, sum_vec);
        float sum = sum_arr[0] + sum_arr[1] + sum_arr[2] + sum_arr[3] + 
                    sum_arr[4] + sum_arr[5] + sum_arr[6] + sum_arr[7];
        for (; j < K; j++) {
            sum += row[j] * x[j];
        }
        out[i] = sum;
    }
}

static void add_bias(float* x, const float* b, int N) {
#pragma omp parallel for
    for (int i = 0; i < N; i++) x[i] += b[i];
}

static void residual_add(float* x, const float* y, int N) {
#pragma omp parallel for
    for (int i = 0; i < N; i++) x[i] += y[i];
}

static void gelu_inplace(float* x, int N) {
    const float c = 0.7978845608f;
#pragma omp parallel for
    for (int i = 0; i < N; i++) {
        float v = x[i];
        float t = tanhf(c * (v + 0.044715f * v * v * v));
        x[i] = 0.5f * v * (1.0f + t);
    }
}

static void softmax_inplace(float* x, int N) {
    float max_val = x[0];
    for (int i = 1; i < N; i++)
        if (x[i] > max_val) max_val = x[i];
    float sum = 0.0f;
    for (int i = 0; i < N; i++) {
        x[i] = expf(x[i] - max_val);
        sum += x[i];
    }
    for (int i = 0; i < N; i++) x[i] /= sum;
}

static void forward(
    int token_id,
    int pos,
    KVCache* kv,
    float* x,
    float* buf,
    float* qkv_buf,
    float* attn_buf,
    float* ff_buf,
    float* logits,
    bool profile_ops = false
) {
    const int C = cfg.n_embd;
    const int H = cfg.n_head;
    const int hs = C / H;

    float* content_row = W.wte + (long long)token_id * C;
    float* pos_row = W.wpe + (long long)pos * C;

#pragma omp parallel for
    for (int i = 0; i < C; i++)
        x[i] = content_row[i] + pos_row[i];

    for (int l = 0; l < cfg.n_layer; l++) {
        auto t_attn_start = std::chrono::high_resolution_clock::now();
        layer_norm(buf, x, W.ln1_w[l], W.ln1_b[l], C);
        matmul_vec(qkv_buf, W.c_attn_w[l], buf, 3 * C, C);
        add_bias(qkv_buf, W.c_attn_b[l], 3 * C);

        float* q = qkv_buf;
        float* k = qkv_buf + C;
        float* v = qkv_buf + 2 * C;

        float* k_cache = kv->k_cache + (long long)l * cfg.block_size * C;
        float* v_cache = kv->v_cache + (long long)l * cfg.block_size * C;

        memcpy(k_cache + (long long)pos * C, k, C * sizeof(float));
        memcpy(v_cache + (long long)pos * C, v, C * sizeof(float));

#pragma omp parallel for
        for (int h = 0; h < H; h++) {
            float* q_h = q + h * hs;
            float scale = 1.0f / sqrtf((float)hs);
            float* local_attn = attn_buf + h * cfg.block_size;

            for (int t = 0; t <= pos; t++) {
                float* k_h = k_cache + (long long)t * C + h * hs;
                float dot = 0.0f;
                for (int d = 0; d < hs; d++)
                    dot += q_h[d] * k_h[d];
                local_attn[t] = dot * scale;
            }

            softmax_inplace(local_attn, pos + 1);

            float* out_h = buf + h * hs;
            memset(out_h, 0, hs * sizeof(float));

            for (int t = 0; t <= pos; t++) {
                float* v_h = v_cache + (long long)t * C + h * hs;
                float a = local_attn[t];
                for (int d = 0; d < hs; d++)
                    out_h[d] += a * v_h[d];
            }
        }

        float* attn_out = qkv_buf;
        matmul_vec(attn_out, W.c_proj_w[l], buf, C, C);
        add_bias(attn_out, W.c_proj_b[l], C);
        residual_add(x, attn_out, C);
        if (profile_ops) {
            auto t_attn_end = std::chrono::high_resolution_clock::now();
            prof_attn_us += std::chrono::duration_cast<std::chrono::microseconds>(t_attn_end - t_attn_start).count();
        }

        auto t_mlp_start = std::chrono::high_resolution_clock::now();
        layer_norm(buf, x, W.ln2_w[l], W.ln2_b[l], C);

        matmul_vec(ff_buf, W.fc_w[l], buf, 4 * C, C);
        add_bias(ff_buf, W.fc_b[l], 4 * C);
        gelu_inplace(ff_buf, 4 * C);

        matmul_vec(buf, W.mlp_proj_w[l], ff_buf, C, 4 * C);
        add_bias(buf, W.mlp_proj_b[l], C);
        residual_add(x, buf, C);
        if (profile_ops) {
            auto t_mlp_end = std::chrono::high_resolution_clock::now();
            prof_mlp_us += std::chrono::duration_cast<std::chrono::microseconds>(t_mlp_end - t_mlp_start).count();
        }
    }

    auto t_lm_start = std::chrono::high_resolution_clock::now();
    layer_norm(buf, x, W.ln_f_w, W.ln_f_b, C);
    matmul_vec(logits, W.lm_head_w, buf, cfg.vocab_size, C);
    if (profile_ops) {
        auto t_lm_end = std::chrono::high_resolution_clock::now();
        prof_lmhead_us += std::chrono::duration_cast<std::chrono::microseconds>(t_lm_end - t_lm_start).count();
    }
}

static void map_weights(float* data, int version) {
    float* ptr = data;
    const int C = cfg.n_embd;
    const int L = cfg.n_layer;

    W.wte = ptr; ptr += (long long)cfg.vocab_size * C;
    W.wpe = ptr; ptr += (long long)cfg.block_size * C;

    W.ln1_w = (float**)malloc(L * sizeof(float*));
    W.ln1_b = (float**)malloc(L * sizeof(float*));
    W.c_attn_w = (float**)malloc(L * sizeof(float*));
    W.c_attn_b = (float**)malloc(L * sizeof(float*));
    W.c_proj_w = (float**)malloc(L * sizeof(float*));
    W.c_proj_b = (float**)malloc(L * sizeof(float*));
    W.ln2_w = (float**)malloc(L * sizeof(float*));
    W.ln2_b = (float**)malloc(L * sizeof(float*));
    W.fc_w = (float**)malloc(L * sizeof(float*));
    W.fc_b = (float**)malloc(L * sizeof(float*));
    W.mlp_proj_w = (float**)malloc(L * sizeof(float*));
    W.mlp_proj_b = (float**)malloc(L * sizeof(float*));

    for (int l = 0; l < L; l++) {
        W.ln1_w[l] = ptr; ptr += C;
        W.ln1_b[l] = ptr; ptr += C;
        W.c_attn_w[l] = ptr; ptr += 3LL * C * C;
        W.c_attn_b[l] = ptr; ptr += 3LL * C;
        W.c_proj_w[l] = ptr; ptr += 1LL * C * C;
        W.c_proj_b[l] = ptr; ptr += C;
        W.ln2_w[l] = ptr; ptr += C;
        W.ln2_b[l] = ptr; ptr += C;
        W.fc_w[l] = ptr; ptr += 4LL * C * C;
        W.fc_b[l] = ptr; ptr += 4LL * C;
        W.mlp_proj_w[l] = ptr; ptr += 1LL * C * 4 * C;
        W.mlp_proj_b[l] = ptr; ptr += C;
    }

    W.ln_f_w = ptr; ptr += C;
    W.ln_f_b = ptr; ptr += C;

    if (version == 2) {
        // Tied weights: lm_head shares wte buffer
        W.lm_head_w = W.wte;
    } else {
        // Untied weights: separate lm_head block
        W.lm_head_w = ptr; ptr += (long long)cfg.vocab_size * C;
    }

    // Print per-tensor table
    fprintf(stderr, "\n--- WEIGHT TENSOR LAYOUT TABLE (Format v%d) ---\n", version);
    fprintf(stderr, "%-26s | %10s | %10s | %12s | %5s\n", "Tensor", "Elements", "Size (MiB)", "File Offset", "Align");
    fprintf(stderr, "----------------------------------------------------------------------\n");
    auto row = [](const char* name, long long elems, void* p) {
        double mib = (double)(elems * sizeof(float)) / (1024.0 * 1024.0);
        uintptr_t offset = (uintptr_t)p - (uintptr_t)model_data_buffer + 64;
        uintptr_t align = (uintptr_t)p % 64;
        fprintf(stderr, "%-26s | %10lld | %10.2f | %12llu | %5llu\n",
                name, elems, mib, (unsigned long long)offset, (unsigned long long)align);
    };

    row("wte", (long long)cfg.vocab_size * C, W.wte);
    row("wpe", (long long)cfg.block_size * C, W.wpe);
    for (int l = 0; l < L; l++) {
        char n[64];
        snprintf(n, sizeof(n), "h.%d.ln1_w", l); row(n, C, W.ln1_w[l]);
        snprintf(n, sizeof(n), "h.%d.ln1_b", l); row(n, C, W.ln1_b[l]);
        snprintf(n, sizeof(n), "h.%d.c_attn_w", l); row(n, 3LL * C * C, W.c_attn_w[l]);
        snprintf(n, sizeof(n), "h.%d.c_attn_b", l); row(n, 3LL * C, W.c_attn_b[l]);
        snprintf(n, sizeof(n), "h.%d.c_proj_w", l); row(n, 1LL * C * C, W.c_proj_w[l]);
        snprintf(n, sizeof(n), "h.%d.c_proj_b", l); row(n, C, W.c_proj_b[l]);
        snprintf(n, sizeof(n), "h.%d.ln2_w", l); row(n, C, W.ln2_w[l]);
        snprintf(n, sizeof(n), "h.%d.ln2_b", l); row(n, C, W.ln2_b[l]);
        snprintf(n, sizeof(n), "h.%d.fc_w", l); row(n, 4LL * C * C, W.fc_w[l]);
        snprintf(n, sizeof(n), "h.%d.fc_b", l); row(n, 4LL * C, W.fc_b[l]);
        snprintf(n, sizeof(n), "h.%d.mlp_proj_w", l); row(n, 1LL * C * 4 * C, W.mlp_proj_w[l]);
        snprintf(n, sizeof(n), "h.%d.mlp_proj_b", l); row(n, C, W.mlp_proj_b[l]);
    }
    row("ln_f_w", C, W.ln_f_w);
    row("ln_f_b", C, W.ln_f_b);
    if (version == 2) {
        row("lm_head_w (TIED->wte)", (long long)cfg.vocab_size * C, W.lm_head_w);
    } else {
        row("lm_head_w (UNTIED)", (long long)cfg.vocab_size * C, W.lm_head_w);
    }
    fprintf(stderr, "----------------------------------------------------------------------\n\n");

    // Assert and verify 64-byte alignment for all weight tensors
    int total_tensors = 0;
    int aligned_tensors = 0;
    auto check_ptr = [&](const char* name, void* p) {
        total_tensors++;
        uintptr_t addr = (uintptr_t)p;
        uintptr_t rem = addr % 64;
        if (rem == 0) aligned_tensors++;
        #ifndef NO_ALIGN_ASSERT
        if (rem != 0) {
            fprintf(stderr, "FATAL: Tensor %s at %p is NOT 64-byte aligned (offset %% 64 = %llu)!\n",
                    name, p, (unsigned long long)rem);
            abort();
        }
        #endif
    };

    check_ptr("wte", W.wte);
    check_ptr("wpe", W.wpe);
    for (int l = 0; l < L; l++) {
        char name[64];
        snprintf(name, sizeof(name), "h.%d.ln1_w", l); check_ptr(name, W.ln1_w[l]);
        snprintf(name, sizeof(name), "h.%d.ln1_b", l); check_ptr(name, W.ln1_b[l]);
        snprintf(name, sizeof(name), "h.%d.c_attn_w", l); check_ptr(name, W.c_attn_w[l]);
        snprintf(name, sizeof(name), "h.%d.c_attn_b", l); check_ptr(name, W.c_attn_b[l]);
        snprintf(name, sizeof(name), "h.%d.c_proj_w", l); check_ptr(name, W.c_proj_w[l]);
        snprintf(name, sizeof(name), "h.%d.c_proj_b", l); check_ptr(name, W.c_proj_b[l]);
        snprintf(name, sizeof(name), "h.%d.ln2_w", l); check_ptr(name, W.ln2_w[l]);
        snprintf(name, sizeof(name), "h.%d.ln2_b", l); check_ptr(name, W.ln2_b[l]);
        snprintf(name, sizeof(name), "h.%d.fc_w", l); check_ptr(name, W.fc_w[l]);
        snprintf(name, sizeof(name), "h.%d.fc_b", l); check_ptr(name, W.fc_b[l]);
        snprintf(name, sizeof(name), "h.%d.mlp_proj_w", l); check_ptr(name, W.mlp_proj_w[l]);
        snprintf(name, sizeof(name), "h.%d.mlp_proj_b", l); check_ptr(name, W.mlp_proj_b[l]);
    }
    check_ptr("ln_f_w", W.ln_f_w);
    check_ptr("ln_f_b", W.ln_f_b);
    check_ptr("lm_head_w", W.lm_head_w);

    fprintf(stderr, "ALIGN_ASSERT: %d/%d weight tensors verified 64-byte aligned (ptr %% 64 == 0)\n",
            aligned_tensors, total_tensors);
}

int main(int argc, char* argv[]) {
    print_mem_stage("process_start");
    auto start_time = std::chrono::high_resolution_clock::now();

    bool dump_logits = false;
    const char* model_path = "model.bin";
    for (int i = 1; i < argc; i++) {
        if (strcmp(argv[i], "--dump-logits") == 0) dump_logits = true;
        if (strcmp(argv[i], "--model") == 0 && i + 1 < argc) model_path = argv[++i];
    }

#ifdef _WIN32
    HANDLE hFile = CreateFileA(model_path, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
    if (hFile == INVALID_HANDLE_VALUE) {
        fprintf(stderr, "ERROR_MODEL_NOT_FOUND: could not open %s\n", model_path);
        return 1;
    }
    ModelHeader header;
    DWORD bytesRead = 0;
    if (!ReadFile(hFile, &header, sizeof(ModelHeader), &bytesRead, NULL) || bytesRead != sizeof(ModelHeader)) {
        fprintf(stderr, "ERROR: Failed to read model header\n");
        CloseHandle(hFile);
        return 1;
    }
    if (memcmp(header.magic, "NANO", 4) != 0) {
        if (*(int*)header.magic == 8) {
            fprintf(stderr, "ERROR: Legacy model format detected (20-byte unaligned header). Expected magic 'NANO' (64-byte aligned header). Model rejected.\n");
        } else {
            fprintf(stderr, "ERROR: Invalid model format. Expected magic 'NANO', got 0x%08X. Model rejected.\n", *(uint32_t*)header.magic);
        }
        CloseHandle(hFile);
        return 1;
    }
    if (header.version != 1 && header.version != 2) {
        fprintf(stderr, "ERROR: Unsupported model version %u. Expected version 1 or 2.\n", header.version);
        CloseHandle(hFile);
        return 1;
    }
    if (header.header_size != 64) {
        fprintf(stderr, "ERROR: Unsupported header size %u. Expected 64 bytes.\n", header.header_size);
        CloseHandle(hFile);
        return 1;
    }
    cfg.n_layer = header.n_layer;
    cfg.n_head = header.n_head;
    cfg.n_embd = header.n_embd;
    cfg.block_size = header.block_size;
    cfg.vocab_size = header.vocab_size;

    HANDLE hMap = CreateFileMappingA(hFile, NULL, PAGE_READONLY, 0, 0, NULL);
    void* mapped_data = MapViewOfFile(hMap, FILE_MAP_READ, 0, 0, 0);
    if (!mapped_data) {
        fprintf(stderr, "ERROR: MapViewOfFile failed\n");
        CloseHandle(hMap);
        CloseHandle(hFile);
        return 1;
    }
    model_data_buffer = (float*)((char*)mapped_data + 64);
#else
    int fd = open(model_path, O_RDONLY);
    if (fd < 0) {
        fprintf(stderr, "ERROR_MODEL_NOT_FOUND: could not open %s\n", model_path);
        return 1;
    }
    struct stat sb;
    if (fstat(fd, &sb) < 0) {
        fprintf(stderr, "ERROR_FSTAT_FAILED\n");
        close(fd);
        return 1;
    }
    ModelHeader header;
    if (read(fd, &header, sizeof(ModelHeader)) != sizeof(ModelHeader)) {
        fprintf(stderr, "ERROR_READING_HEADER\n");
        close(fd);
        return 1;
    }
    if (memcmp(header.magic, "NANO", 4) != 0) {
        if (*(int*)header.magic == 8) {
            fprintf(stderr, "ERROR: Legacy model format detected (20-byte unaligned header). Expected magic 'NANO' (64-byte aligned header). Model rejected.\n");
        } else {
            fprintf(stderr, "ERROR: Invalid model format. Expected magic 'NANO', got 0x%08X. Model rejected.\n", *(uint32_t*)header.magic);
        }
        close(fd);
        return 1;
    }
    if (header.version != 1 && header.version != 2) {
        fprintf(stderr, "ERROR: Unsupported model version %u. Expected version 1 or 2.\n", header.version);
        close(fd);
        return 1;
    }
    if (header.header_size != 64) {
        fprintf(stderr, "ERROR: Unsupported header size %u. Expected 64 bytes.\n", header.header_size);
        close(fd);
        return 1;
    }
    cfg.n_layer = header.n_layer;
    cfg.n_head = header.n_head;
    cfg.n_embd = header.n_embd;
    cfg.block_size = header.block_size;
    cfg.vocab_size = header.vocab_size;

    void* mapped_data = mmap(NULL, sb.st_size, PROT_READ, MAP_SHARED, fd, 0);
    if (mapped_data == MAP_FAILED) {
        fprintf(stderr, "ERROR_MMAP_FAILED\n");
        close(fd);
        return 1;
    }
    model_data_buffer = (float*)((char*)mapped_data + 64);
#endif

    print_mem_stage("after_map");

    map_weights(model_data_buffer, header.version);

    // Touch all weights to measure working set when ALL mapped weight pages are in RAM
    volatile float sum_touch = 0.0f;
    long long total_floats = (header.version == 2) ? 51082752LL : 76814336LL;
    for (long long i = 0; i < total_floats; i += 1024) {
        sum_touch += model_data_buffer[i];
    }
    print_mem_stage("after_touching_weights");

    const int C = cfg.n_embd;
    const int L = cfg.n_layer;

    // Allocate all buffers ONCE
    KVCache kv;
    kv.k_cache = (float*)malloc((long long)L * cfg.block_size * C * sizeof(float));
    kv.v_cache = (float*)malloc((long long)L * cfg.block_size * C * sizeof(float));

    float* x = (float*)malloc(C * sizeof(float));
    float* buf = (float*)malloc(C * sizeof(float));
    float* qkv_buf = (float*)malloc(3 * C * sizeof(float));
    float* attn_buf = (float*)malloc(cfg.n_head * cfg.block_size * sizeof(float));
    float* ff_buf = (float*)malloc(4 * C * sizeof(float));
    float* logits = (float*)malloc(cfg.vocab_size * sizeof(float));

    print_mem_stage("after_kv_activation_alloc");

    auto model_loaded_time = std::chrono::high_resolution_clock::now();
    long long load_ms = std::chrono::duration_cast<std::chrono::milliseconds>(model_loaded_time - start_time).count();
    fprintf(stderr, "ENGINE_READY: load_ms=%lld\n", load_ms);
    fflush(stderr);

    char line[8192];
    bool first_token_measured = false;

    while (fgets(line, sizeof(line), stdin)) {
        if (strncmp(line, "GEN ", 4) != 0) continue;

        char req_id[64];
        int max_new_tokens, top_k;
        unsigned int seed;
        float temperature;

        char* token_str = line + 4;
        sscanf(token_str, "%s %d %f %d %u", req_id, &max_new_tokens, &temperature, &top_k, &seed);

        int spaces = 0;
        char* ids_str = token_str;
        while (*ids_str && spaces < 5) {
            if (*ids_str == ' ') spaces++;
            ids_str++;
        }

        std::vector<int> input_ids;
        char* tok = strtok(ids_str, ", \n");
        while (tok) {
            input_ids.push_back(atoi(tok));
            tok = strtok(NULL, ", \n");
        }

        if (input_ids.size() >= (size_t)cfg.block_size)
            input_ids.resize(cfg.block_size - 1);

        if (input_ids.size() + max_new_tokens > (size_t)cfg.block_size) {
            max_new_tokens = cfg.block_size - input_ids.size();
        }

        if (temperature < 0.01f) temperature = 0.01f;
        if (top_k < 1) top_k = 1;
        if (top_k > cfg.vocab_size) top_k = cfg.vocab_size;

        srand(seed);

        // Reset per-op profiling timers
        prof_attn_us = 0;
        prof_mlp_us = 0;
        prof_lmhead_us = 0;
        prof_sample_us = 0;

        auto req_start = std::chrono::high_resolution_clock::now();

        FILE* f_logits = NULL;
        if (dump_logits) f_logits = fopen("logits.bin", "wb");

        for (int i = 0; i < (int)input_ids.size(); i++) {
            forward(input_ids[i], i, &kv, x, buf, qkv_buf, attn_buf, ff_buf, logits, false);
            if (f_logits) fwrite(logits, sizeof(float), cfg.vocab_size, f_logits);
        }
        if (f_logits) fclose(f_logits);

        auto prefill_done_time = std::chrono::high_resolution_clock::now();
        long long prefill_ms = std::chrono::duration_cast<std::chrono::milliseconds>(prefill_done_time - req_start).count();

        int pos = input_ids.size();
        int generated = 0;

        for (int i = 0; i < max_new_tokens; i++) {
            if (pos >= cfg.block_size) break;

            auto t_sample_start = std::chrono::high_resolution_clock::now();
            for (int v = 0; v < cfg.vocab_size; v++) logits[v] /= temperature;

            std::vector<std::pair<float, int>> pairs(cfg.vocab_size);
            for (int v = 0; v < cfg.vocab_size; v++) pairs[v] = {logits[v], v};

            std::partial_sort(pairs.begin(), pairs.begin() + top_k, pairs.end(),
                [](const std::pair<float,int>& a, const std::pair<float,int>& b) {
                    return a.first > b.first;
                });

            float sum = 0.0f;
            for (int j = 0; j < top_k; j++) {
                pairs[j].first = expf(pairs[j].first);
                sum += pairs[j].first;
            }
            for (int j = 0; j < top_k; j++) pairs[j].first /= sum;

            float r = (float)rand() / ((float)RAND_MAX + 1.0f);
            float cum = 0.0f;
            int next_id = pairs[0].second;
            for (int j = 0; j < top_k; j++) {
                cum += pairs[j].first;
                if (r <= cum) {
                    next_id = pairs[j].second;
                    break;
                }
            }
            auto t_sample_end = std::chrono::high_resolution_clock::now();
            prof_sample_us += std::chrono::duration_cast<std::chrono::microseconds>(t_sample_end - t_sample_start).count();

            printf("T %d\n", next_id);
            fflush(stdout);

            if (!first_token_measured) {
                print_mem_stage("after_first_token");
                first_token_measured = true;
            }

            generated++;
            if (next_id == 50256) break; 

            // Decode forward step with profiling
            forward(next_id, pos, &kv, x, buf, qkv_buf, attn_buf, ff_buf, logits, true);
            pos++;
        }

        auto decode_done_time = std::chrono::high_resolution_clock::now();
        long long decode_ms = std::chrono::duration_cast<std::chrono::milliseconds>(decode_done_time - prefill_done_time).count();

        // Print per-op profiling summary to stderr
        long long prof_total = prof_attn_us + prof_mlp_us + prof_lmhead_us + prof_sample_us;
        if (prof_total > 0) {
            fprintf(stderr, "OP_PROFILE: attn=%lldus (%.1f%%) mlp=%lldus (%.1f%%) lm_head=%lldus (%.1f%%) sample=%lldus (%.1f%%) total=%lldus\n",
                    prof_attn_us, (double)prof_attn_us * 100.0 / prof_total,
                    prof_mlp_us, (double)prof_mlp_us * 100.0 / prof_total,
                    prof_lmhead_us, (double)prof_lmhead_us * 100.0 / prof_total,
                    prof_sample_us, (double)prof_sample_us * 100.0 / prof_total,
                    prof_total);
            fflush(stderr);
        }

        printf("DONE %lld %lld %d\n", prefill_ms, decode_ms, generated);
        fflush(stdout);
    }

    return 0;
}