#!/usr/bin/env python3
"""Shared helpers for the three fingerprint benchmarks (LLMmap / LLM Verify /
llm-fingerprint-detector).

These benchmarks do not go through EvalScope's dataset pipeline. Each runner
script probes the target OpenAI-compatible endpoint with its own tool logic and
writes a report JSON shaped like EvalScope reports:

    output/<folder>/<benchmark>/seed_<seed>/reports/<benchmark>.json

with at least ``score`` (float 0~1) and ``num`` so that
``bash/collect_results.py`` can aggregate them like any other benchmark.
"""

import argparse
import json
import urllib.error
import urllib.request

# 各工具报告文件名与其 benchmark 名一致
BENCHMARK_LLMMAP = 'llmmap'
BENCHMARK_LLM_VERIFY = 'llm_verify'
BENCHMARK_DETECTOR = 'llm_fingerprint_detector'

ALL_FINGERPRINT_BENCHMARKS = [BENCHMARK_LLMMAP, BENCHMARK_LLM_VERIFY, BENCHMARK_DETECTOR]

# 默认对被测端点关闭 thinking：指纹探测需要稳定的可见回答，
# 思考链会烧掉 max_tokens 且改变输出分布。sglang/vLLM 均支持该字段。
DEFAULT_EXTRA_BODY = {'chat_template_kwargs': {'thinking': False}}


def add_common_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
    """CLI arguments shared by all three fingerprint runners."""
    parser.add_argument('--api-url', required=True,
                        help='Target OpenAI-compatible API base URL, e.g. http://localhost:30000/v1')
    parser.add_argument('--model', required=True, help='Served model name to probe')
    parser.add_argument('--report-path', required=True,
                        help='Where to write the EvalScope-style report JSON')
    parser.add_argument('--timeout', type=int, default=120,
                        help='Per-request timeout in seconds (default: %(default)s)')
    parser.add_argument('--thinking', action='store_true', default=False,
                        help='Do NOT disable thinking on the target (default: disabled)')
    return parser


def chat_completion(api_url: str, model: str, user_prompt: str,
                    system_prompt: str = '', temperature: float = 1.0,
                    max_tokens: int = 512, timeout: int = 120,
                    extra_body: dict = None):
    """Minimal OpenAI chat-completions call (stdlib only).

    Returns:
        (content, error) — exactly one of them is None.
    """
    messages = []
    if system_prompt:
        messages.append({'role': 'system', 'content': system_prompt})
    messages.append({'role': 'user', 'content': user_prompt})

    payload = {
        'model': model,
        'messages': messages,
        'temperature': temperature,
        'max_tokens': max_tokens,
        'stream': False,
    }
    payload.update(extra_body or {})

    req = urllib.request.Request(
        f"{api_url.rstrip('/')}/chat/completions",
        data=json.dumps(payload).encode('utf-8'),
        headers={'Content-Type': 'application/json'},
        method='POST',
    )
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            data = json.loads(resp.read().decode('utf-8'))
    except urllib.error.HTTPError as e:
        detail = ''
        try:
            detail = e.read().decode('utf-8')[:200]
        except Exception:
            pass
        return None, f'HTTP {e.code}: {detail}'
    except Exception as e:
        return None, f'request failed: {e}'

    choices = data.get('choices') or []
    if not choices:
        return None, 'empty choices in response'
    message = choices[0].get('message') or {}
    content = message.get('content')
    # 部分推理模型把可见内容放在 reasoning_content；仅当 content 为空时兜底。
    if not content:
        content = message.get('reasoning_content') or ''
    return str(content), None


def write_report(report_path: str, benchmark: str, score: float, num: int,
                 **details) -> None:
    """Write an EvalScope-style report JSON consumable by collect_results.py."""
    from pathlib import Path

    report_path = Path(report_path)
    report_path.parent.mkdir(parents=True, exist_ok=True)
    payload = {
        'benchmark': benchmark,
        'score': float(score),
        'num': int(num),
    }
    payload.update(details)
    report_path.write_text(json.dumps(payload, ensure_ascii=False, indent=2),
                           encoding='utf-8')
    print(f'[fingerprint] report written: {report_path}')