#!/usr/bin/env python3
"""LLMmap fingerprint benchmark runner.

把目标端点当作"未知模型"：向其发送 LLMmap 的 8 条指纹查询，收集回答后用
LLMmap 预训练 open-set 模型与 52 个已知模板比对，输出 Top-K 及得分。

必须用装好 torch/transformers 的解释器运行（默认 llmmap conda 环境），
由 run.py 以子进程方式调用：

    <llmmap-python> run_llmmap.py --api-url ... --model ... --report-path ...

得分（score ∈ [0,1]）：
    - 提供 --expected-model 时：Top-1 模板与期望模型名匹配 → 1.0，否则 0.0
      （匹配为归一化后的包含关系，如 "GLM-5.2" 可匹配 "zai-org/GLM-5.2"）。
    - 未提供时：置信度 score = max(0, 1 - top1_distance / --distance-scale)。
"""

import argparse
import os
import sys

# 嵌入模型已缓存到本地，禁止联网检查更新
os.environ.setdefault('HF_HUB_OFFLINE', '1')
os.environ.setdefault('TRANSFORMERS_OFFLINE', '1')

from common import BENCHMARK_LLMMAP, add_common_args, chat_completion, write_report


def normalize_name(name: str) -> str:
    """小写并去掉组织前缀/斜杠/冒号后的空白，便于宽松匹配。"""
    n = str(name).strip().lower()
    if '/' in n:
        n = n.split('/')[-1]
    return n.replace('-', '').replace('_', '').replace('.', '')


def main():
    parser = argparse.ArgumentParser(description='LLMmap fingerprint benchmark')
    add_common_args(parser)
    parser.add_argument('--tools-root', default='/data1/xii',
                        help='Directory containing the cloned LLMmap repo (default: %(default)s)')
    parser.add_argument('--llmmap-model-path', default=None,
                        help='Pretrained LLMmap open-set model directory '
                             '(default: <tools-root>/LLMmap/data/pretrained_models/default)')
    parser.add_argument('--device', default='cpu', choices=['cpu', 'cuda'])
    parser.add_argument('--temperature', type=float, default=0.7,
                        help='Sampling temperature when querying the target (default: %(default)s)')
    parser.add_argument('--max-tokens', type=int, default=512,
                        help='Max tokens per target answer (default: %(default)s)')
    parser.add_argument('--expected-model', default=None,
                        help='Ground-truth model identity; when set, score is a strict match flag')
    parser.add_argument('--distance-scale', type=float, default=60.0,
                        help='Confidence normalizer when no expected model is given '
                             '(observed: same-family ~20, others ~40+)')
    parser.add_argument('-k', type=int, default=5, help='Top-K templates to record')
    args = parser.parse_args()

    llmmap_root = os.path.join(args.tools_root, 'LLMmap')
    if not os.path.isdir(llmmap_root):
        print(f'ERROR: LLMmap repo not found at {llmmap_root}')
        sys.exit(1)
    model_path = args.llmmap_model_path or os.path.join(
        llmmap_root, 'data', 'pretrained_models', 'default')
    sys.path.insert(0, llmmap_root)

    from LLMmap.inference import load_LLMmap

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

    # 逐条向被测端点发送指纹查询
    extra_body = None if args.thinking else {'chat_template_kwargs': {'thinking': False}}
    answers, errors = [], []
    for i, query in enumerate(llmmap.queries, 1):
        content, err = chat_completion(
            args.api_url, args.model, query,
            temperature=args.temperature, max_tokens=args.max_tokens,
            timeout=args.timeout, extra_body=extra_body,
        )
        if err:
            print(f'  query {i}/{len(llmmap.queries)} failed: {err}')
            errors.append({'query_index': i - 1, 'error': err})
            content = ''
        else:
            print(f'  query {i}/{len(llmmap.queries)} ok ({len(content)} chars)')
        answers.append(content or '')

    # 与已知模板比对（open-set 距离检索）
    # 端点大面积失败时回答为空，距离毫无意义 —— 直接判失败而不是给假分数
    n_ok = len(answers) - len(errors)
    if n_ok <= len(answers) // 2:
        write_report(
            args.report_path, BENCHMARK_LLMMAP, 0.0,
            num=len(answers),
            score_mode='error',
            top1=None,
            topk=[],
            expected_model=args.expected_model,
            n_query_errors=len(errors),
            query_errors=errors[:5],
            error=f'too many failed queries ({len(errors)}/{len(answers)}); '
                  f'is the endpoint up and serving --model?',
        )
        print(f'[llmmap] FAILED: {len(errors)}/{len(answers)} queries errored')
        sys.exit(1)

    distances = llmmap(answers)
    order = sorted(range(len(distances)), key=lambda i: distances[i])
    label_map = llmmap.label_map  # {index: template_name}
    topk = [{'name': label_map[i], 'distance': float(distances[i])}
            for i in order[:max(1, args.k)]]

    top1_name, top1_dist = topk[0]['name'], topk[0]['distance']
    if args.expected_model:
        matched = normalize_name(args.expected_model) in normalize_name(top1_name) or \
                  normalize_name(top1_name) in normalize_name(args.expected_model)
        score = 1.0 if matched else 0.0
        score_mode = 'identity_match'
    else:
        score = max(0.0, 1.0 - float(top1_dist) / args.distance_scale)
        score_mode = 'confidence'

    write_report(
        args.report_path, BENCHMARK_LLMMAP, score,
        num=len(answers),
        score_mode=score_mode,
        top1=topk[0],
        topk=topk,
        expected_model=args.expected_model,
        n_query_errors=len(errors),
        query_errors=errors[:5],
    )
    print(f"[llmmap] Top-1: {top1_name} (distance={top1_dist:.4f}) -> score={score:.4f}")


if __name__ == '__main__':
    main()