// PRE Native Inference Engine — llama.cpp C API integration
// Compiled as C++ (required by llama.cpp), exposes C-linkage API via inference.h.

#include "inference.h"
#include <llama.h>

#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cmath>
#include <string>
#include <vector>
#include <sys/time.h>
#include <sys/stat.h>
#include <unistd.h>
#include <pwd.h>

// ── Helpers ─────────────────────────────────────────────────────────────────

static double now_ms() {
    struct timeval tv;
    gettimeofday(&tv, NULL);
    return tv.tv_sec * 1000.0 + tv.tv_usec / 1000.0;
}

// ── Engine struct ───────────────────────────────────────────────────────────

struct inference_engine {
    struct llama_model   *model;
    struct llama_context *ctx;
    struct llama_sampler *sampler;
    const struct llama_vocab *vocab;
    int n_ctx;
    int n_kv_used; // tokens consumed in KV cache
    char model_path[1024];
};

// ── Init / Free ─────────────────────────────────────────────────────────────

extern "C" inference_engine_t *inference_init(const char *model_path, int n_ctx) {
    if (!model_path || access(model_path, R_OK) != 0) {
        fprintf(stderr, "[inference] Model file not found: %s\n", model_path ? model_path : "(null)");
        return NULL;
    }

    fprintf(stderr, "[inference] Loading model: %s\n", model_path);
    double t0 = now_ms();

    // Model params: offload all layers to GPU, use mmap
    struct llama_model_params mparams = llama_model_default_params();
    mparams.n_gpu_layers = -1; // offload all layers
    mparams.use_mmap = true;

    // Progress callback
    mparams.progress_callback = [](float progress, void *) -> bool {
        static int last_pct = -1;
        int pct = (int)(progress * 100);
        if (pct != last_pct && pct % 10 == 0) {
            fprintf(stderr, "[inference] Loading... %d%%\n", pct);
            last_pct = pct;
        }
        return true;
    };

    struct llama_model *model = llama_model_load_from_file(model_path, mparams);
    if (!model) {
        fprintf(stderr, "[inference] Failed to load model\n");
        return NULL;
    }

    double load_ms = now_ms() - t0;
    fprintf(stderr, "[inference] Model loaded in %.1fs\n", load_ms / 1000.0);

    // Context params
    struct llama_context_params cparams = llama_context_default_params();
    cparams.n_ctx = n_ctx > 0 ? (uint32_t)n_ctx : 131072;
    cparams.n_batch = 2048;
    cparams.n_ubatch = 512;
    cparams.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED;
    cparams.offload_kqv = true;
    cparams.n_threads = 4; // Apple Silicon benefits from fewer threads (GPU does most work)
    cparams.n_threads_batch = 8;

    struct llama_context *ctx = llama_init_from_model(model, cparams);
    if (!ctx) {
        fprintf(stderr, "[inference] Failed to create context\n");
        llama_model_free(model);
        return NULL;
    }

    // Sampler chain: top_k -> min_p -> top_p -> temp -> dist
    // Matches PRE's default sampling: temperature=1.0, top_k=64, top_p=0.95, min_p=0.05
    struct llama_sampler_chain_params sparams = llama_sampler_chain_default_params();
    struct llama_sampler *sampler = llama_sampler_chain_init(sparams);
    llama_sampler_chain_add(sampler, llama_sampler_init_top_k(64));
    llama_sampler_chain_add(sampler, llama_sampler_init_min_p(0.05f, 1));
    llama_sampler_chain_add(sampler, llama_sampler_init_top_p(0.95f, 1));
    llama_sampler_chain_add(sampler, llama_sampler_init_temp(1.0f));
    llama_sampler_chain_add(sampler, llama_sampler_init_dist(LLAMA_DEFAULT_SEED));

    inference_engine_t *eng = (inference_engine_t *)calloc(1, sizeof(inference_engine_t));
    eng->model = model;
    eng->ctx = ctx;
    eng->sampler = sampler;
    eng->vocab = llama_model_get_vocab(model);
    eng->n_ctx = (int)llama_n_ctx(ctx);
    eng->n_kv_used = 0;
    strncpy(eng->model_path, model_path, sizeof(eng->model_path) - 1);

    fprintf(stderr, "[inference] Native inference initialized (ctx=%d)\n", eng->n_ctx);
    return eng;
}

extern "C" void inference_free(inference_engine_t *eng) {
    if (!eng) return;
    if (eng->sampler) llama_sampler_free(eng->sampler);
    if (eng->ctx) llama_free(eng->ctx);
    if (eng->model) llama_model_free(eng->model);
    free(eng);
}

// ── Generate ────────────────────────────────────────────────────────────────

extern "C" char *inference_generate(inference_engine_t *eng,
                                     const inference_message_t *messages, int n_messages,
                                     int max_tokens,
                                     inference_token_cb on_token, void *userdata,
                                     inference_stats_t *stats_out) {
    if (!eng || !messages || n_messages <= 0) return NULL;
    if (max_tokens <= 0) max_tokens = 16384;

    // Step 1: Apply chat template to format the prompt
    // Build llama_chat_message array
    std::vector<llama_chat_message> chat(n_messages);
    for (int i = 0; i < n_messages; i++) {
        chat[i].role = messages[i].role;
        chat[i].content = messages[i].content;
    }

    // Get template from model metadata (NULL = auto-detect from GGUF)
    int32_t prompt_len = llama_chat_apply_template(
        NULL, chat.data(), chat.size(), true, NULL, 0);
    if (prompt_len < 0) {
        fprintf(stderr, "[inference] Chat template not supported by model\n");
        return NULL;
    }

    std::string prompt(prompt_len + 1, '\0');
    llama_chat_apply_template(NULL, chat.data(), chat.size(), true,
                              &prompt[0], prompt_len + 1);
    prompt.resize(prompt_len);

    // Step 2: Tokenize
    int n_prompt_max = prompt_len + 256;
    std::vector<llama_token> tokens(n_prompt_max);
    int n_tokens = llama_tokenize(eng->vocab, prompt.c_str(), prompt_len,
                                   tokens.data(), n_prompt_max, true, true);
    if (n_tokens < 0) {
        // Buffer too small, resize and retry
        n_prompt_max = -n_tokens + 64;
        tokens.resize(n_prompt_max);
        n_tokens = llama_tokenize(eng->vocab, prompt.c_str(), prompt_len,
                                   tokens.data(), n_prompt_max, true, true);
        if (n_tokens < 0) {
            fprintf(stderr, "[inference] Tokenization failed\n");
            return NULL;
        }
    }
    tokens.resize(n_tokens);

    // Check if prompt fits in context
    if (n_tokens + max_tokens > eng->n_ctx) {
        fprintf(stderr, "[inference] Prompt (%d tokens) + max_tokens (%d) exceeds context (%d)\n",
                n_tokens, max_tokens, eng->n_ctx);
        // Reduce max_tokens to fit
        max_tokens = eng->n_ctx - n_tokens - 64;
        if (max_tokens < 128) {
            fprintf(stderr, "[inference] Not enough context space\n");
            return NULL;
        }
    }

    // Step 3: Check KV cache prefix match
    // If the beginning of the new prompt matches what's already in the KV cache,
    // we can skip those tokens (prompt caching).
    // For simplicity in this initial implementation, we clear the cache each time.
    // TODO: implement proper prefix caching by comparing token sequences
    llama_memory_clear(llama_get_memory(eng->ctx), true);
    eng->n_kv_used = 0;

    // Step 4: Prefill — process all prompt tokens
    double t_prefill_start = now_ms();

    // Process in batches
    int n_processed = 0;
    while (n_processed < n_tokens) {
        int n_batch = std::min(2048, n_tokens - n_processed);
        struct llama_batch batch = llama_batch_get_one(
            tokens.data() + n_processed, n_batch);
        if (llama_decode(eng->ctx, batch) != 0) {
            fprintf(stderr, "[inference] Decode failed during prefill at token %d\n", n_processed);
            return NULL;
        }
        n_processed += n_batch;
    }

    double t_prefill_ms = now_ms() - t_prefill_start;
    eng->n_kv_used = n_tokens;

    // Step 5: Generate tokens
    double t_gen_start = now_ms();
    std::string output;
    output.reserve(max_tokens * 4); // rough estimate

    char piece_buf[256];
    int n_generated = 0;

    for (int i = 0; i < max_tokens; i++) {
        // Sample the next token
        llama_token new_token = llama_sampler_sample(eng->sampler, eng->ctx, -1);

        // Check for end of generation
        if (llama_vocab_is_eog(eng->vocab, new_token)) {
            break;
        }

        // Convert token to text
        int n_chars = llama_token_to_piece(eng->vocab, new_token,
                                            piece_buf, sizeof(piece_buf) - 1, 0, true);
        if (n_chars < 0) {
            // Should not happen with valid tokens
            continue;
        }
        piece_buf[n_chars] = '\0';

        // Append to output
        output.append(piece_buf, n_chars);
        n_generated++;

        // Callback for streaming
        if (on_token) {
            if (!on_token(piece_buf, n_chars, userdata)) {
                break; // Caller requested abort
            }
        }

        // Decode the new token to update KV cache
        struct llama_batch batch = llama_batch_get_one(&new_token, 1);
        if (llama_decode(eng->ctx, batch) != 0) {
            fprintf(stderr, "[inference] Decode failed at generated token %d\n", n_generated);
            break;
        }
    }

    double t_gen_ms = now_ms() - t_gen_start;
    eng->n_kv_used += n_generated;

    // Fill stats
    if (stats_out) {
        stats_out->prompt_tokens = n_tokens;
        stats_out->generated_tokens = n_generated;
        stats_out->prompt_eval_ms = t_prefill_ms;
        stats_out->generate_ms = t_gen_ms;
        stats_out->tokens_per_sec = t_gen_ms > 0 ? (n_generated * 1000.0 / t_gen_ms) : 0;
        stats_out->prompt_tok_per_sec = t_prefill_ms > 0 ? (n_tokens * 1000.0 / t_prefill_ms) : 0;
    }

    // Return a copy of the output
    char *result = (char *)malloc(output.size() + 1);
    memcpy(result, output.c_str(), output.size());
    result[output.size()] = '\0';
    return result;
}

// ── KV Cache Persistence ────────────────────────────────────────────────────

extern "C" bool inference_save_kv(inference_engine_t *eng, const char *path) {
    if (!eng || !path) return false;
    // llama_state_save_file saves the full context state including KV cache
    size_t n_token_count = 0;
    bool ok = llama_state_save_file(eng->ctx, path, NULL, n_token_count);
    if (ok) {
        fprintf(stderr, "[inference] KV cache saved to %s\n", path);
    } else {
        fprintf(stderr, "[inference] Failed to save KV cache to %s\n", path);
    }
    return ok;
}

extern "C" bool inference_load_kv(inference_engine_t *eng, const char *path) {
    if (!eng || !path) return false;
    struct stat st;
    if (stat(path, &st) != 0) return false;

    size_t n_token_count = 0;
    bool ok = llama_state_load_file(eng->ctx, path, NULL, 0, &n_token_count);
    if (ok) {
        eng->n_kv_used = (int)n_token_count;
        fprintf(stderr, "[inference] KV cache loaded from %s (%zu tokens)\n", path, n_token_count);
    } else {
        fprintf(stderr, "[inference] Failed to load KV cache from %s\n", path);
    }
    return ok;
}

extern "C" int inference_kv_used(inference_engine_t *eng) {
    return eng ? eng->n_kv_used : 0;
}

extern "C" void inference_kv_clear(inference_engine_t *eng) {
    if (!eng) return;
    llama_memory_clear(llama_get_memory(eng->ctx), true);
    eng->n_kv_used = 0;
}

// ── Ollama Model Resolver ───────────────────────────────────────────────────

extern "C" char *inference_resolve_ollama_model(const char *tag) {
    if (!tag) return NULL;

    // Parse tag into name and version (e.g., "pre-gemma4" -> "pre-gemma4:latest")
    std::string name(tag);
    std::string version = "latest";
    size_t colon = name.find(':');
    if (colon != std::string::npos) {
        version = name.substr(colon + 1);
        name = name.substr(0, colon);
    }

    // Build manifest path: ~/.ollama/models/manifests/registry.ollama.ai/library/<name>/<version>
    const char *home = getenv("HOME");
    if (!home) {
        struct passwd *pw = getpwuid(getuid());
        home = pw ? pw->pw_dir : "/tmp";
    }

    char manifest_path[2048];
    snprintf(manifest_path, sizeof(manifest_path),
             "%s/.ollama/models/manifests/registry.ollama.ai/library/%s/%s",
             home, name.c_str(), version.c_str());

    FILE *f = fopen(manifest_path, "r");
    if (!f) {
        fprintf(stderr, "[inference] Ollama manifest not found: %s\n", manifest_path);
        return NULL;
    }

    // Read manifest JSON
    fseek(f, 0, SEEK_END);
    long fsize = ftell(f);
    fseek(f, 0, SEEK_SET);
    std::string json(fsize, '\0');
    fread(&json[0], 1, fsize, f);
    fclose(f);

    // Find the model layer (mediaType contains "image.model")
    // Simple parser — look for "digest":"sha256:..." after "image.model"
    size_t model_pos = json.find("image.model");
    if (model_pos == std::string::npos) {
        fprintf(stderr, "[inference] No model layer in manifest\n");
        return NULL;
    }

    size_t digest_pos = json.find("\"digest\":", model_pos);
    if (digest_pos == std::string::npos) {
        // Try before the model_pos (digest comes before mediaType in some formats)
        digest_pos = json.rfind("\"digest\":", model_pos);
    }
    if (digest_pos == std::string::npos) {
        fprintf(stderr, "[inference] No digest found for model layer\n");
        return NULL;
    }

    // Extract digest value: "sha256:hexstring"
    size_t val_start = json.find('"', digest_pos + 9);
    size_t val_end = json.find('"', val_start + 1);
    if (val_start == std::string::npos || val_end == std::string::npos) return NULL;

    std::string digest = json.substr(val_start + 1, val_end - val_start - 1);
    // Convert "sha256:abc123" to "sha256-abc123" for Ollama blob path
    size_t colon2 = digest.find(':');
    if (colon2 != std::string::npos) digest[colon2] = '-';

    // Build blob path
    char blob_path[2048];
    snprintf(blob_path, sizeof(blob_path), "%s/.ollama/models/blobs/%s",
             home, digest.c_str());

    // Verify it exists
    if (access(blob_path, R_OK) != 0) {
        fprintf(stderr, "[inference] GGUF blob not found: %s\n", blob_path);
        return NULL;
    }

    return strdup(blob_path);
}
