// dual_ctx.cpp — valide Path A: prefill HTP (ctx ngl99) -> transfert KV -> decode CPU (ctx ngl0).
// Mesure: RAM des 2 loads, débit prefill HTP, débit decode CPU, cohérence de la continuation.
// Build: NDK clang++ -Iinclude, link dist/lib. Run: GGML_HEXAGON_GDN_PREFILL=1 LD_LIBRARY_PATH=lib ./dual_ctx model.gguf
#include <cstdio>
#include <cstring>
#include <cstdint>
#include <cstdlib>
#include <string>
#include <vector>
#include <ctime>
#include "llama.h"
#include "ggml-backend.h"

static int64_t usec(){ struct timespec ts; clock_gettime(CLOCK_MONOTONIC,&ts); return (int64_t)ts.tv_sec*1000000+ts.tv_nsec/1000; }
static long rss_mb(){ FILE*f=fopen("/proc/self/status","r"); if(!f)return -1; char l[256]; long k=0;
  while(fgets(l,sizeof l,f)) if(sscanf(l,"VmRSS: %ld kB",&k)==1) break; fclose(f); return k/1024; }

int main(int argc, char** argv){
  if(argc<2){ printf("usage: dual_ctx model.gguf\n"); return 1; }
  llama_backend_init();
  fprintf(stderr,"RSS start            : %ld MB\n", rss_mb());

  const char* e=getenv("KZ_NGL_H"); int ngl_h = e? atoi(e):99;  // control: KZ_NGL_H=0 => ctx_h aussi CPU
  auto mp_h=llama_model_default_params(); mp_h.n_gpu_layers=ngl_h;
  // sélection explicite du device HTP0 (comme -dev HTP0) pour ctx_h
  static ggml_backend_dev_t devs[2]={nullptr,nullptr};
  for(size_t i=0;i<ggml_backend_dev_count();i++){ auto d=ggml_backend_dev_get(i);
    fprintf(stderr,"dev[%zu]=%s\n", i, ggml_backend_dev_name(d));
    if(!strcmp(ggml_backend_dev_name(d),"HTP0")) devs[0]=d; }
  if(ngl_h>0 && devs[0]){ mp_h.devices=devs; fprintf(stderr,"ctx_h: device HTP0 explicite\n"); }
  auto m_h=llama_model_load_from_file(argv[1],mp_h); if(!m_h){printf("nomodel htp\n");return 1;}
  fprintf(stderr,"RSS after HTP(ngl99) : %ld MB\n", rss_mb());
  auto mp_c=llama_model_default_params(); mp_c.n_gpu_layers=0;
  auto m_c=llama_model_load_from_file(argv[1],mp_c); if(!m_c){printf("nomodel cpu\n");return 1;}
  fprintf(stderr,"RSS after +CPU(ngl0) : %ld MB  (delta = coût 2e instance)\n", rss_mb());

  auto base=llama_context_default_params(); base.n_ctx=2048; base.n_batch=2048;
  base.flash_attn_type=LLAMA_FLASH_ATTN_TYPE_ENABLED; base.type_k=GGML_TYPE_F16; base.type_v=GGML_TYPE_F16;
  auto cph=base; cph.n_threads=8; auto ctx_h=llama_init_from_model(m_h,cph);
  auto cpc=base; cpc.n_threads=4; auto ctx_c=llama_init_from_model(m_c,cpc);
  auto vo=llama_model_get_vocab(m_h);
  auto s=llama_sampler_init_greedy();
  fprintf(stderr,"RSS after 2 ctx      : %ld MB\n", rss_mb());

  int rep = (argc>2)? atoi(argv[2]) : 1;  // argv[2]=N: répète la phrase (gonfle le prompt pour test perf)
  std::string pr="<|im_start|>system\nTu es Kazeia, soutien psy bref en francais, tutoiement.<|im_end|>\n<|im_start|>user\n";
  for(int i=0;i<rep;i++) pr+="Je rumine la nuit et je n'arrive plus a dormir depuis des semaines, je me sens vide. ";
  pr+="<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n";
  int n=-llama_tokenize(vo,pr.c_str(),pr.size(),0,0,true,true); std::vector<llama_token> t(n);
  llama_tokenize(vo,pr.c_str(),pr.size(),t.data(),n,true,true);
  fprintf(stderr,"prompt tokens        : %d\n", n);

  int64_t a=usec(); auto b=llama_batch_get_one(t.data(),n);
  if(llama_decode(ctx_h,b)!=0){printf("prefill fail\n");return 1;}
  int64_t bb=usec(); fprintf(stderr,"PREFILL HTP          : %d tok / %.0f ms = %.1f t/s\n", n,(bb-a)/1e3, n*1e6/(double)(bb-a));

  size_t sz=llama_state_seq_get_size(ctx_h,0); std::vector<uint8_t> buf(sz);
  llama_state_seq_get_data(ctx_h, buf.data(), sz, 0);
  size_t got=llama_state_seq_set_data(ctx_c, buf.data(), sz, 0);
  fprintf(stderr,"KV transfer HTP->CPU : %zu bytes (set=%zu)\n", sz, got);

  // REF: decode quelques tokens sur le contexte HTP lui-même (sans transfert) -> isole prefill vs transfert
  { llama_token id=llama_sampler_sample(s,ctx_h,-1); std::string ref; int p=n;
    for(int i=0;i<12;i++){ char z[128]; int l=llama_token_to_piece(vo,id,z,sizeof z,0,true); if(l>0)ref.append(z,l);
      if(llama_vocab_is_eog(vo,id))break; llama_token tk=id; llama_pos pp=p; int32_t ns=1; llama_seq_id sd=0,*sp=&sd; int8_t lg=1;
      llama_batch sb; memset(&sb,0,sizeof sb); sb.n_tokens=1; sb.token=&tk; sb.pos=&pp; sb.n_seq_id=&ns; sb.seq_id=&sp; sb.logits=&lg;
      if(llama_decode(ctx_h,sb)!=0)break; p++; id=llama_sampler_sample(s,ctx_h,-1); }
    fprintf(stderr,"REF decode sur ctx HTP (memes poids, pas de transfert): '%s'\n", ref.c_str()); }
  // restaure le snapshot dans ctx_h (le self-decode l'a avancé) pour repartir propre
  llama_state_seq_set_data(ctx_h, buf.data(), sz, 0);

  llama_token id=llama_sampler_sample(s,ctx_h,-1);
  fprintf(stderr,"first id (from HTP prefill logits): %d\n", id);
  std::string out; int pos=n, ndec=0; int64_t d0=usec();
  for(int i=0;i<64;i++){
    char z[128]; int l=llama_token_to_piece(vo,id,z,sizeof z,0,true); if(l>0) out.append(z,l);
    if(i<8) fprintf(stderr,"  tok[%d]=%d piece='%.*s' eog=%d\n", i, id, l>0?l:0, l>0?z:"", llama_vocab_is_eog(vo,id));
    if(llama_vocab_is_eog(vo,id)) break;
    llama_token tok=id; llama_pos p=pos; int32_t nsi=1; llama_seq_id sid=0, *sidp=&sid; int8_t lg=1;
    llama_batch sb; memset(&sb,0,sizeof sb);
    sb.n_tokens=1; sb.token=&tok; sb.pos=&p; sb.n_seq_id=&nsi; sb.seq_id=&sidp; sb.logits=&lg;
    if(llama_decode(ctx_c,sb)!=0){fprintf(stderr,"decode fail @%d\n",i);break;}
    pos++; ndec++; id=llama_sampler_sample(s,ctx_c,-1);
  }
  int64_t d1=usec();
  fprintf(stderr,"DECODE CPU           : %d tok = %.1f t/s\n", ndec, ndec*1e6/(double)(d1-d0));
  fprintf(stderr,"RSS peak             : %ld MB\n", rss_mb());
  printf("OUT: %s\n", out.c_str());
  return 0;
}