#include "llama.h"
#include "common.h" // Hilfsfunktionen aus dem llama.cpp Repository
#include <iostream>
#include <string>
#include <vector>

// GBNF Grammatik-Definition, die das Modell auf exakte JSON-Tool-Calls zwingt
const std::string JSON_GRAMMAR = 
    "root   ::= object\n"
    "object ::= \"{\n  \\\"action\\\": \\\"\" name \"\\\",\n  \\\"parameters\\\": \" params \"\n}\"\n"
    "name   ::= [a-zA-Z0-9_]+\n"
    "params ::= \"{\n    \" pair ( \",\n    \" pair )* \"\n  }\"\n"
    "pair   ::= \"\\\"\" [a-zA-Z0-9_]+ \"\\\": \\\"\" [^\"]* \"\\\"\"";

int main(int argc, char ** argv) {
    // 1. Hardware-Parameter für Smartphone-SoCs initialisieren
    llama_model_params model_params = llama_model_default_params();
    
    #if defined(__ARM_NEON) || defined(__APPLE__)
        model_params.use_mmap = true;  // Direktes Mapping in den RAM (schont die Batterie)
        model_params.main_gpu = 0;     // Nutze GPU/NPU falls via Metal/Vulkan kompiliert
    #endif

    // 2. Das kompakte 300MB Modell laden
    std::string model_path = "slm750_edge_q4_k.gguf";
    llama_model * model = llama_load_model_from_file(model_path.c_str(), model_params);
    if (!model) {
        std::cerr << "Fehler: Modell konnte nicht geladen werden!" << std::endl;
        return 1;
    }

    // 3. Kontext für hocheffizientes Batch-Processing aufsetzen
    llama_context_params ctx_params = llama_context_default_params();
    ctx_params.n_ctx   = 2048; // Gleiche Kontext-Länge wie im Trainings-Skript von Neo
    ctx_params.n_batch = 512;
    ctx_params.n_threads = 4; // Optimal für die "Big-Cores" moderner Smartphone-CPUs

    llama_context * ctx = llama_new_context_with_model(model, ctx_params);

    // 4. Grammatik-Parser initialisieren (Der "Wow-Effekt" Garant für 100% JSON-Validität)
    auto grammar = llama_grammar_init(JSON_GRAMMAR.c_str());

    // 5. Prompt für den Tool-Call vorbereiten
    std::string prompt = "User: Erstelle eine neue Stellenausschreibung für einen Senior App Developer.\nAgent:";
    std::vector<llama_token> tokens = llama_tokenize(ctx, prompt, true);

    // Prompt in den KV-Cache evaluieren
    llama_decode(ctx, llama_batch_get_one(tokens.data(), tokens.size(), 0, 0));

    std::cout << "Generiere ultraschnellen Tool-Call...\n" << std::endl;

    // 6. Token-Generierungs-Schleife mit Grammatik-Zwang
    for (int i = 0; i < 200; ++i) {
        auto   candidates = llama_sampler_get_candidates(ctx);
        
        // Zwinge das Modell, nur Tokens zu wählen, die zur JSON-Grammatik passen
        llama_grammar_sample(grammar, ctx, &candidates);
        
        llama_token token = llama_sampler_sample(ctx, &candidates);

        // Abbruch bei End-of-Token
        if (token == llama_token_eos(model)) break;

        // Token sofort auf dem Bildschirm ausgeben
        std::cout << llama_token_to_piece(ctx, token);
        std::flush(std::cout);

        // Nächstes Token verarbeiten
        llama_decode(ctx, llama_batch_get_one(&token, 1, i + tokens.size(), 0));
    }

    // Speicherbereinigung
    llama_grammar_free(grammar);
    llama_free(ctx);
    llama_free_model(model);
    return 0;
}