// EPQ target with Qwen3.6 MTP draft context.
#include "llama.h"
#include "sdq_pager.h"
#include "common.h"
#include "sampling.h"
#include "speculative.h"

#include <algorithm>
#include <chrono>
#include <cstdio>
#include <cstdlib>
#include <string>
#include <vector>

using namespace sdq;

namespace {
struct opt_t { std::string model, prompt; int ctx=4096, threads=8, n=48, draft_n=2, ram=8192, ubatch=128; bool no_pf=false; };
bool parse(int argc, char ** argv, opt_t & o) {
    auto next = [&](int & i) -> const char * { return ++i < argc ? argv[i] : nullptr; };
    for (int i=1; i<argc; i++) {
        const std::string a=argv[i]; const char * v=nullptr;
        if (a=="-m" || a=="--model") { if (!(v=next(i))) return false; o.model=v; }
        else if (a=="-p" || a=="--prompt") { if (!(v=next(i))) return false; o.prompt=v; }
        else if (a=="-n") { if (!(v=next(i))) return false; o.n=atoi(v); }
        else if (a=="--draft-n") { if (!(v=next(i))) return false; o.draft_n=atoi(v); }
        else if (a=="-c") { if (!(v=next(i))) return false; o.ctx=atoi(v); }
        else if (a=="-t") { if (!(v=next(i))) return false; o.threads=atoi(v); }
        else if (a=="-b") { if (!(v=next(i))) return false; o.ubatch=atoi(v); }
        else if (a=="--ram-budget-mb") { if (!(v=next(i))) return false; o.ram=atoi(v); }
        else if (a=="--no-prefetch") o.no_pf=true;
        else return false;
    }
    return !o.model.empty() && !o.prompt.empty();
}
double now_s() { return std::chrono::duration<double>(std::chrono::steady_clock::now().time_since_epoch()).count(); }
}

int main(int argc, char ** argv) {
    opt_t o;
    if (!parse(argc, argv, o)) return 1;
    setenv("SDQ_MODEL_PATH", o.model.c_str(), 1);
    setenv("SDQ_IO_THREADS", std::to_string(o.threads*3).c_str(), 1);
    setenv("SDQ_RAM_BUDGET_MB", std::to_string(o.ram).c_str(), 1);
    setenv("SDQ_FADV_DONTNEED", "1", 1);
    if (o.no_pf) setenv("SDQ_PREFETCH", "0", 1);
    llama_backend_init();
    llama_numa_init(GGML_NUMA_STRATEGY_DISABLED);
    llama_model_params mp=llama_model_default_params();
    mp.n_gpu_layers=0; mp.load_mode=LLAMA_LOAD_MODE_MMAP; mp.load_mtp=true;
    sdq_set_enabled(true); sdq_install_buft_overrides(&mp);
    llama_model * model=llama_model_load_from_file(o.model.c_str(), mp);
    if (!model || !sdq_init(o.model)) return 2;
    llama_context_params c=llama_context_default_params();
    c.n_ctx=o.ctx; c.n_seq_max=1; c.n_batch=std::max(o.ubatch,o.draft_n+1); c.n_ubatch=o.ubatch;
    c.n_outputs_max=o.draft_n+1; c.n_outputs_max_per_seq=o.draft_n+1;
    c.n_threads=o.threads; c.n_threads_batch=o.threads; c.ctx_type=LLAMA_CONTEXT_TYPE_DEFAULT;
    llama_context * ctx=llama_init_from_model(model,c);
    if (!ctx) return 3;

    common_params p;
    p.model.path=o.model; p.n_ctx=o.ctx; p.n_batch=c.n_batch; p.n_ubatch=o.ubatch;
    p.cpuparams.n_threads=o.threads; p.cpuparams_batch.n_threads=o.threads;
    p.speculative.types={COMMON_SPECULATIVE_TYPE_DRAFT_MTP};
    p.speculative.draft.n_max=o.draft_n;
    auto init=common_speculative_init_from_params(p,model,ctx);
    llama_context * dctx=init->context();
    p.speculative.draft.ctx_tgt=ctx; p.speculative.draft.ctx_dft=dctx;
    if (!dctx) { fprintf(stderr,"[sdq-mtp] MTP context init failed\n"); return 4; }
    common_speculative * spec=common_speculative_init(p.speculative,1);
    if (!spec) { fprintf(stderr,"[sdq-mtp] speculative init failed\n"); return 5; }
    common_params_sampling sp; sp.temp=0; sp.top_k=1;
    common_sampler * sampler=common_sampler_init(model,sp);
    auto input=common_tokenize(ctx,o.prompt,true,true);
    if (input.empty()) return 6;
    common_batch pb(ctx);
    for (size_t i=0;i+1<input.size();i++) pb.add(input[i],(llama_pos)i,0,false);
    if (pb.size() && llama_process(ctx,LLAMA_PROCESS_TYPE_DECODE,pb.get())!=0) return 7;
    pager::instance().end_token(); common_speculative_process(spec,pb);
    llama_tokens prompt(input.begin(),input.end()-1); common_speculative_begin(spec,0,prompt);
    llama_token last=input.back(); int past=(int)prompt.size(), gen=0, drafted=0, accepted=0;
    common_prompt_checkpoint ck; llama_tokens draft;
    const double t0=now_s(); printf("assistant: "); fflush(stdout);
    while (gen<o.n) {
        ck.update_pos(prompt.size(),llama_memory_seq_pos_min(llama_get_memory(ctx),0),llama_memory_seq_pos_max(llama_get_memory(ctx),0));
        ck.update_tgt(ctx,0,LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); ck.update_dft(dctx,0,LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
        if (draft.empty()) { auto & dp=common_speculative_get_draft_params(spec,0); dp={}; dp.drafting=true; dp.n_max=std::min(o.draft_n,o.n-gen-1); dp.pos0=past; dp.id_last=last; dp.prompt=&prompt; dp.result=&draft; common_speculative_draft(spec); }
        ck.load_dft(dctx,0,LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
        common_batch vb(ctx); vb.add(last,past++,0,true); for(size_t i=0;i<draft.size();i++) vb.add(draft[i],past+(llama_pos)i,0,true);
        if (llama_process(ctx,LLAMA_PROCESS_TYPE_DECODE,vb.get())!=0) break;
        pager::instance().end_token();
        if(!common_speculative_process(spec,vb)) break;
        auto ids=common_sampler_sample_and_accept_n(sampler,ctx,draft); if(ids.empty()) break;
        drafted+=(int)draft.size(); accepted+=(int)ids.size()-1;
        if(ids.size()-1<draft.size()) { ck.load_tgt(ctx,0,LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); ck.load_dft(dctx,0,LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); prompt.resize(ck.n_tokens); past=ck.n_tokens; draft=std::move(ids); continue; }
        for(auto id:ids) { if(gen>=o.n) break; printf("%s",common_token_to_piece(ctx,id).c_str()); gen++; prompt.push_back(last); last=id; }
        common_speculative_accept(spec,0,(uint16_t)(ids.size()-1)); past+=(int)ids.size()-1;
    }
    printf("\n"); double dt=now_s(); fprintf(stderr,"[sdq-mtp] generated=%d drafted=%d accepted=%d acceptance=%.1f%% speed=%.2f tok/s\n",gen,drafted,accepted,drafted?100.0*accepted/drafted:0,dt-t0?gen/(dt-t0):0);
    common_speculative_print_stats(spec); common_speculative_free(spec); common_sampler_free(sampler); sdq_shutdown(); llama_free(ctx); llama_model_free(model); llama_backend_free(); return 0;
}