#!/usr/bin/env python3
"""FP-Fusion extra cells 采集（10 个新增 cell：binary-* + day-of-week × en/zh）。
对齐 fp_fusion engine D 层采样；带 enable_thinking=false 禁思考 + api-key。"""
import argparse
import asyncio
import json
import math
import random
import sys
import time
from collections import Counter
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))  # 仓库根/site-packages, 使 evalharness 包可导入
import httpx
from evalharness.fingerprint.battery import NEW_CELLS, D_TEMPERATURE
from evalharness.fingerprint.engine import normalize_answer

SYS_ONE_WORD = {'en': 'Answer with exactly one word. No punctuation, no explanation.',
                'zh': '只回答一个词，不要标点，不要解释。'}
client_key = ''


def normalize_for_ref(cell_id, raw):
    domains = {f"{c['id']}:{lang}": c['domain'] for c in NEW_CELLS for lang in ('en', 'zh')}
    return normalize_answer(raw, domains[cell_id])


async def sample_cell(client, api_url, model, cell_id, pool, n, sem, retries, timeout):
    samples = []
    for i in range(n):
        prompt = random.choice(pool)
        body = {'model': model,
                'messages': [{'role': 'system',
                              'content': SYS_ONE_WORD['zh' if cell_id.endswith(':zh') else 'en']},
                             {'role': 'user', 'content': prompt}],
                'temperature': D_TEMPERATURE, 'max_tokens': 16, 'stream': False,
                'chat_template_kwargs': {'enable_thinking': False}}
        headers = {'Content-Type': 'application/json'}
        if client_key:
            headers['Authorization'] = f'Bearer {client_key}'
        for attempt in range(retries):
            try:
                async with sem:
                    r = await client.post(f"{api_url.rstrip('/')}/chat/completions",
                                          json=body, headers=headers, timeout=timeout)
                    r.raise_for_status()
                    data = r.json()
                content = (data.get('choices') or [{}])[0].get('message', {}).get('content') or ''
                samples.append(content)
                break
            except Exception as e:
                if attempt == retries - 1:
                    print(f'  [{cell_id}] sample {i} failed: {str(e)[:80]}', file=sys.stderr)
                    samples.append('')
                else:
                    await asyncio.sleep(1.5 * (attempt + 1))
    return cell_id, samples


async def main_async(args):
    sem = asyncio.Semaphore(args.concurrency)
    timeout = httpx.Timeout(max(args.timeout, 120))
    async with httpx.AsyncClient(timeout=timeout) as client:
        tasks = []
        for c in NEW_CELLS:
            for lang in ('en', 'zh'):
                cell_id = f"{c['id']}:{lang}"
                tasks.append(asyncio.create_task(sample_cell(
                    client, args.api_url, args.model, cell_id, c['par'][lang],
                    args.samples, sem, args.retries, timeout)))
        results = await asyncio.gather(*tasks)
    cells = {}
    for cell_id, samples in results:
        counts = Counter()
        valid = invalid = refusal = empty = error = 0
        for s in samples:
            if s == '':
                error += 1; continue
            norm, cat = normalize_for_ref(cell_id, s)
            if cat == 'valid' and norm is not None:
                counts[norm] += 1; valid += 1
            elif cat == 'refusal': refusal += 1
            elif cat == 'empty': empty += 1
            else: invalid += 1
        total = len(samples)
        entropy = 0.0
        if total and counts:
            entropy = -sum((v / total) * math.log2(v / total) for v in counts.values())
        cells[cell_id] = {'cellId': cell_id, 'counts': {str(k): v for k, v in counts.items()},
                          'validCount': valid, 'invalidCount': invalid,
                          'refusalCount': refusal, 'emptyCount': empty,
                          'errorCount': error, 'totalCount': total,
                          'entropyBits': entropy, 'normalizedEntropy': 0.0,
                          'medianLatencyMs': None, 'meanCompletionTokens': None,
                          'meanReasoningTokens': None}
        print(f'  {cell_id}: valid={valid}/{total}', flush=True)
    return cells


def main():
    p = argparse.ArgumentParser(description='Collect fp_fusion extra reference cells')
    p.add_argument('--api-url', required=True)
    p.add_argument('--model', required=True)
    p.add_argument('--api-key', default='')
    p.add_argument('--out', required=True)
    p.add_argument('--samples', type=int, default=25)
    p.add_argument('--concurrency', type=int, default=2)
    p.add_argument('--retries', type=int, default=8)
    p.add_argument('--timeout', type=float, default=120)
    args = p.parse_args()

    global client_key
    client_key = args.api_key

    cells = asyncio.run(main_async(args))
    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    payload = {'formatVersion': 1, 'protocol': 'one-token/v1', 'model': args.model,
               'collectedAt': time.strftime('%Y-%m-%dT%H:%M:%S.000Z', time.gmtime()),
               'samplesPerCell': args.samples, 'postReasoning': False,
               'extraCellsOnly': True, 'cells': cells}
    out.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding='utf-8')
    print(f'extra reference written -> {out} ({len(cells)} cells)')


if __name__ == '__main__':
    main()