// Persistent CPU classifier. Each D request is one exact HF token window.
// State is cleared before EVERY window. This process never opens datasets.
#include "llama.h"
#include "ggml-cpu.h"
#include "qwen35-specialized.h"
#include <chrono>
#include <cmath>
#include <iomanip>
#include <iostream>
#include <sstream>
#include <stdexcept>
#include <string>
#include <vector>
#include <cstdlib>
#include <cstring>
using Clock = std::chrono::steady_clock;
static double seconds(Clock::time_point t) { return std::chrono::duration<double>(Clock::now()-t).count(); }

extern "C" void q35_clear_packed_weights();
extern "C" size_t q35_packed_weight_bytes();
struct Runtime {
    llama_model *model=nullptr;
    llama_context *ctx=nullptr;
    ggml_threadpool_t pool=nullptr;
    llama_batch batch{};
    bool allocated=false;
    const char *path;
    int threads;
    Runtime(const char *p,int t):path(p),threads(t) {}
    void unload() {
        if(allocated) { llama_batch_free(batch); allocated=false; }
        if(ctx) { llama_detach_threadpool(ctx);llama_free(ctx);ctx=nullptr; }
        if(pool) {ggml_threadpool_free(pool);pool=nullptr;}
        if(model) { q35_clear_packed_weights();llama_model_free(model); model=nullptr; }
    }
    ~Runtime() { unload(); }
    void load() {
        if(model) throw std::runtime_error("already loaded");
        auto start=Clock::now();
        auto mp=llama_model_default_params();
        mp.n_gpu_layers=0;mp.use_mmap=true;mp.check_tensors=true;
        model=llama_model_load_from_file(path,mp);
        if(!model || llama_model_n_cls_out(model)!=14) throw std::runtime_error("classifier model load/dimension failure");
        char architecture[64] = {};
        llama_model_meta_val_str(model, "general.architecture", architecture, sizeof(architecture));
        if (std::strcmp(architecture, "qwen35") || llama_model_n_layer(model)!=4 ||
            llama_model_n_embd(model)!=1024)
            throw std::runtime_error("this runtime supports only Qwen3.5 4L / 1024-wide / 14-label classification");
        double weights=seconds(start);start=Clock::now();
        auto cp=llama_context_default_params();
        cp.n_ctx=256;cp.n_batch=256;cp.n_ubatch=256;cp.n_outputs_max=256;cp.n_seq_max=1;
        cp.n_threads=threads;cp.n_threads_batch=threads;cp.embeddings=true;
        cp.pooling_type=LLAMA_POOLING_TYPE_RANK;cp.attention_type=LLAMA_ATTENTION_TYPE_CAUSAL;
        cp.flash_attn_type=LLAMA_FLASH_ATTN_TYPE_DISABLED;cp.type_k=GGML_TYPE_F32;cp.type_v=GGML_TYPE_F32;
        cp.offload_kqv=false;cp.op_offload=false;
        ctx=llama_init_from_model(model,cp);
        if(!ctx) throw std::runtime_error("context load failure");
        if(q35_specialize_level()>=9) {
            auto pp=ggml_threadpool_params_default(threads);
            pp.poll=0; // Sleep between requests; no idle CPU spinning.
            pool=ggml_threadpool_new(&pp);
            if(!pool) throw std::runtime_error("thread pool allocation failure");
            llama_attach_threadpool(ctx,pool,pool);
        }
        batch=llama_batch_init(256,0,1);allocated=true;
        std::cout<<"{\"event\":\"ready\",\"weights_seconds\":"<<weights
                 <<",\"context_seconds\":"<<seconds(start)<<",\"cache_dtype\":\"F32\",\"gpu_layers\":0}\n"<<std::flush;
    }
    void predict(const std::vector<int>& ids) {
        if(!ctx) throw std::runtime_error("not loaded");
        const int vocab=llama_vocab_n_tokens(llama_model_get_vocab(model));
        for(int v:ids) if(v<0 || v>=vocab) throw std::runtime_error("invalid token ID");
        batch.n_tokens=static_cast<int>(ids.size());
        for(int i=0;i<batch.n_tokens;++i) {
            batch.token[i]=ids[i];batch.pos[i]=i;batch.n_seq_id[i]=1;batch.seq_id[i][0]=0;batch.logits[i]=1;
        }
        const auto start=Clock::now();
        llama_memory_clear(llama_get_memory(ctx),true);
        if(llama_decode(ctx,batch)!=0) throw std::runtime_error("decode failure");
        const float *logits=llama_get_embeddings_seq(ctx,0);
        if(!logits) throw std::runtime_error("missing classifier output");
        for(int i=0;i<14;++i) if(!std::isfinite(logits[i])) throw std::runtime_error("nonfinite classifier output");
        double compute=seconds(start);
        std::cout<<"{\"event\":\"prediction\",\"compute_seconds\":"<<compute<<",\"tokens\":[";
        for(size_t i=0;i<ids.size();++i) std::cout<<(i?",":"")<<ids[i];
        std::cout<<"],\"logits\":[";
        for(int i=0;i<14;++i) std::cout<<(i?",":"")<<logits[i];
        std::cout<<"],\"packed_weight_bytes\":"<<q35_packed_weight_bytes()<<"}\n"<<std::flush;
    }
};

int main(int argc,char **argv) {
    bool initialized=false;
    try {
        if(argc!=4) throw std::runtime_error("usage: classifier MODEL.gguf LIBDIR THREADS");
        setenv("SCAMGUARDIAN_QWEN35_NO_KV", "1", 1);
        setenv("SCAMGUARDIAN_QWEN35_NO_KV_ALLOC", "1", 1);
        int threads=std::stoi(argv[3]);if(threads<1 || threads>64) throw std::runtime_error("invalid threads");
        std::cout<<std::setprecision(10);
        auto start=Clock::now();ggml_backend_load_all_from_path(argv[2]);llama_backend_init();initialized=true;
        std::cout<<"{\"event\":\"runtime\",\"seconds\":"<<seconds(start)<<"}\n"<<std::flush;
        {
            Runtime runtime(argv[1],threads);runtime.load();
            std::string line;
            while(std::getline(std::cin,line)) {
                if(line=="QUIT") break;
                if(line=="UNLOAD") {runtime.unload();std::cout<<"{\"event\":\"unloaded\"}\n"<<std::flush;continue;}
                if(line=="LOAD") {runtime.load();continue;}
                try {
                    if(line.size()>8192) throw std::runtime_error("oversized request");
                    std::istringstream input(line);std::string command;int length;
                    if(!(input>>command>>length) || command!="D" || length<1 || length>256) throw std::runtime_error("invalid request");
                    std::vector<int> ids(length);for(auto &id:ids) if(!(input>>id)) throw std::runtime_error("truncated request");
                    std::string extra;if(input>>extra) throw std::runtime_error("trailing request input");
                    runtime.predict(ids);
                } catch(const std::exception&) {
                    // Avoid returning input text. The next request still clears all state.
                    std::cout<<"{\"event\":\"request_error\"}\n"<<std::flush;
                }
            }
        }
        llama_backend_free();return 0;
    } catch(const std::exception& error) {
        std::cerr<<"native classifier error: "<<error.what()<<"\n";
        if(initialized) llama_backend_free();
        return 1;
    }
}