// dual_ctx_mt.cpp — valide l'option C en MULTI-TOUR: par tour, re-prefill full history sur HTP
// (+OPFILTER=SSM_CONV) -> transfert KV -> decode CPU. Vérifie cohérence, mémoire conversationnelle, perf.
// Run: GGML_HEXAGON_OPFILTER=SSM_CONV GGML_HEXAGON_GDN_PREFILL=1 LD_LIBRARY_PATH=lib ./dual_ctx_mt 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; }

struct Msg { const char* role; std::string text; };

static std::string build_prompt(const char* sys, const std::vector<Msg>& hist, const std::string& user){
  std::string p = "<|im_start|>system\n"; p+=sys; p+="<|im_end|>\n";
  for(auto& m: hist){ p+="<|im_start|>"; p+=m.role; p+="\n"; p+=m.text; p+="<|im_end|>\n"; }
  p += "<|im_start|>user\n"; p+=user; p+="<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n";
  return p;
}

int main(int argc, char** argv){
  if(argc<2){ printf("usage: dual_ctx_mt model.gguf\n"); return 1; }
  llama_backend_init();

  // device HTP0 explicite 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);
    if(!strcmp(ggml_backend_dev_name(d),"HTP0")) devs[0]=d; }
  auto mp_h=llama_model_default_params(); mp_h.n_gpu_layers=99; if(devs[0]) mp_h.devices=devs;
  auto m_h=llama_model_load_from_file(argv[1],mp_h); if(!m_h){printf("nomodel htp\n");return 1;}
  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;}

  auto base=llama_context_default_params(); base.n_ctx=4096; base.n_batch=4096;
  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 après 2 modèles+2 ctx : %ld MB\n", rss_mb());

  const char* sys="Tu es Kazeia, soutien psy bienveillant, francais, tutoiement, 2-3 phrases.";
  const char* users[3]={
    "Je m'appelle Marc, j'ai 42 ans, et je dors tres mal depuis trois semaines.",
    "Qu'est-ce que je pourrais essayer ce soir pour mieux dormir ?",
    "Rappelle-moi : quel est mon prenom et mon age, selon ce que je t'ai dit ?" };

  std::vector<Msg> hist;
  for(int k=0;k<3;k++){
    std::string prompt = build_prompt(sys, hist, users[k]);
    int n=-llama_tokenize(vo,prompt.c_str(),prompt.size(),0,0,true,true);
    std::vector<llama_token> t(n); llama_tokenize(vo,prompt.c_str(),prompt.size(),t.data(),n,true,true);

    // re-prefill full history sur HTP (KV vidé)
    llama_memory_clear(llama_get_memory(ctx_h), true);
    int64_t a=usec(); auto b=llama_batch_get_one(t.data(),n);
    if(llama_decode(ctx_h,b)!=0){printf("prefill fail tour %d\n",k);return 1;}
    int64_t bb=usec(); double pf=n*1e6/(double)(bb-a);

    // transfert KV HTP -> CPU
    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);
    llama_memory_clear(llama_get_memory(ctx_c), true);
    llama_state_seq_set_data(ctx_c, buf.data(), sz, 0);

    // decode sur CPU
    llama_token id=llama_sampler_sample(s,ctx_h,-1);
    std::string resp; int pos=n, ndec=0; int64_t d0=usec();
    for(int i=0;i<100;i++){
      char z[128]; int l=llama_token_to_piece(vo,id,z,sizeof z,0,true);
      if(llama_vocab_is_eog(vo,id)) break;
      if(l>0) resp.append(z,l);
      llama_token tk=id; llama_pos p=pos; 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=&p; sb.n_seq_id=&ns; sb.seq_id=&sp; sb.logits=&lg;
      if(llama_decode(ctx_c,sb)!=0){fprintf(stderr,"decode fail t%d@%d\n",k,i);break;}
      pos++; ndec++; id=llama_sampler_sample(s,ctx_c,-1);
    }
    int64_t d1=usec(); double dec=ndec*1e6/(double)(d1-d0);
    printf("\n===== TOUR %d (prompt %d tok | prefill HTP %.1f t/s | decode CPU %.1f t/s) =====\n", k+1, n, pf, dec);
    printf("USER : %s\n", users[k]);
    printf("KAZEIA: %s\n", resp.c_str());
    hist.push_back({"user",users[k]}); hist.push_back({"assistant",resp});
  }
  fprintf(stderr,"RSS peak : %ld MB\n", rss_mb());
  return 0;
}