#!/usr/bin/env python
"""Non-interactive LLM fingerprinting helper.

Collect answers from a real target LLM for the 8 fingerprinting queries,
put one answer per line in a text file, then run:

    python run_identify.py answers.txt [-k 6] [--model_path ./data/pretrained_models/default]
"""
import argparse
import os

from LLMmap.inference import load_LLMmap


def main():
    ap = argparse.ArgumentParser(description='Fingerprint an LLM from a file of answers')
    ap.add_argument('answers_file', type=str, help='Text file with one answer per line (8 lines)')
    ap.add_argument('-k', type=int, default=6, help='Number of top candidates to print')
    ap.add_argument('--model_path', type=str, default='./data/pretrained_models/default')
    ap.add_argument('--device', type=str, default='cpu', choices=['cpu', 'cuda'])
    ap.add_argument('--dump-queries', action='store_true',
                    help='Only print the fingerprinting queries, then exit')
    args = ap.parse_args()

    conf, llmmap = load_LLMmap(args.model_path, device=args.device)

    if args.dump_queries:
        print('Send these queries to the target LLM (one at a time) and save each response '
              'on its own line in your answers file:\n')
        for i, q in enumerate(llmmap.queries):
            print(f'[{i + 1}] {q}\n')
        return

    with open(args.answers_file) as f:
        answers = [line.rstrip('\n') for line in f if line.strip() != '']

    if len(answers) != len(llmmap.queries):
        raise SystemExit(
            f'Expected {len(llmmap.queries)} answers (one per fingerprinting query), '
            f'got {len(answers)}. Use --dump-queries to list the queries.'
        )

    print('### Predicted identity ###')
    llmmap.print_result(llmmap(answers), k=args.k)


if __name__ == '__main__':
    main()