#!/usr/bin/env python
"""Fingerprint a REAL local HF model living outside the 52-template DB (open-set demo).

Steps:
  1. Load a local model (e.g. /data1/models/Qwen3-4B) directly from disk.
  2. Run the 8 LLMmap fingerprinting queries against it and collect real answers.
  3. Feed those answers to the LLMmap open-set predictor and print the top-K
     nearest known templates.
"""
import argparse
import os

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

from LLMmap.inference import load_LLMmap


def load_llm_generator(model_dir, device_map='auto', torch_dtype=torch.bfloat16, max_new_tokens=100):
    tok = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token
    model = AutoModelForCausalLM.from_pretrained(
        model_dir,
        torch_dtype=torch_dtype,
        device_map=device_map,
        trust_remote_code=True,
    )
    model.eval()

    def generate(query, thinking=False):
        messages = [{'role': 'user', 'content': query}]
        # Qwen3 chat templates accept an enable_thinking flag if the tokenizer supports it
        kwargs = {}
        if 'enable_thinking' in tok.chat_template:
            kwargs['enable_thinking'] = thinking
        prompt = tok.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True, **kwargs)
        in_toks = tok(prompt, return_tensors='pt', add_special_tokens=False,
                      return_token_type_ids=False).to(model.device)
        with torch.no_grad():
            out_toks = model.generate(
                **in_toks,
                max_new_tokens=max_new_tokens,
                do_sample=False,
                pad_token_id=tok.eos_token_id,
                eos_token_id=tok.eos_token_id,
            )
        gen = [out_toks[i, in_toks.input_ids.shape[1]:] for i in range(len(out_toks))]
        return tok.batch_decode(gen, skip_special_tokens=True)[0]

    return generate, tok


def main():
    ap = argparse.ArgumentParser(description='Fingerprint a real local HF model (open-set)')
    ap.add_argument('--model_dir', required=True,
                    help='Path to a local HuggingFace model, e.g. /data1/models/Qwen3-4B')
    ap.add_argument('--model_path', default='./data/pretrained_models/default')
    ap.add_argument('--save_answers', default='/tmp/real_answers.txt', help='Where to store answers')
    ap.add_argument('--max_new_tokens', type=int, default=100)
    ap.add_argument('--k', type=int, default=6)
    ap.add_argument('--device', default='cpu', choices=['cpu', 'cuda'])
    args = ap.parse_args()

    model_name = os.path.basename(args.model_dir.rstrip('/'))
    print(f'[1/3] Loading local model: {args.model_dir}')
    generate, tok = load_llm_generator(args.model_dir, torch_dtype=torch.bfloat16,
                                       max_new_tokens=args.max_new_tokens)

    print(f'[2/3] Loading LLMmap predictor: {args.model_path}')
    conf, llmmap = load_LLMmap(args.model_path, device=args.device)

    print(f'[3/3] Running {len(llmmap.queries)} fingerprinting queries against {model_name}...')
    answers = []
    for i, q in enumerate(llmmap.queries, 1):
        print(f'  -- query {i}/{len(llmmap.queries)}')
        try:
            ans = generate(q)
        except Exception as e:
            ans = f'[generation error: {e}]'
        answers.append(ans)
        print(f'     -> {ans[:160].replace(chr(10), " ")}')

    with open(args.save_answers, 'w') as f:
        f.write('\n'.join(answers))
    print(f'\nAnswers saved to {args.save_answers}\n')

    print('### LLMmap open-set prediction (target = real %s) ###' % model_name)
    llmmap.print_result(llmmap(answers), k=args.k)


if __name__ == '__main__':
    main()