"""Load the public panel (JevlinAI/supersearch-tests) into a compact pickle: per call, the ranked canonical URLs."""
import gzip, json, os, pickle, sys, collections
HERE = os.path.dirname(os.path.abspath(__file__))
D = os.environ.get('JEVLIN_PANEL_DIR', os.path.join(HERE, 'panel'))  # a checkout of the supersearch-tests dataset
sys.path.insert(0, D)
import score  # canonicalisation and matching rules, unchanged from the dataset
ARMS = ['exa', 'exa-lite', 'brave', 'parallel-search', 'perplexity-search', 'fast', 'balanced']
items = {r['id']: r for r in map(json.loads, open(os.path.join(D, 'items.jsonl')))}
calls = collections.defaultdict(dict)
for arm in ARMS:
    with gzip.open(os.path.join(D, 'calls', arm + '.jsonl.gz'), 'rt') as f:
        for l in f:
            r = json.loads(l)
            m = r['response'].get('merged') or []
            calls[arm][(r['id'], r['phrasing'], r['trial'])] = dict(
                urls=[x.get('url') or '' for x in m][:20],
                found_by=[[fb.get('provider') if isinstance(fb, dict) else fb for fb in (x.get('found_by') or [])] for x in m][:20],
                votes=[x.get('votes') for x in m][:20],
                latency_ms=r.get('latency_ms'), first_ok=r.get('first_attempt_ok'),
                cost=(r['response'].get('usage') or {}).get('cost_micros'),
                cache=(r['response'].get('usage') or {}).get('cache_hit'),
                router_ms=(r['response'].get('usage') or {}).get('latency_ms'))
pickle.dump((items, dict(calls)), open(os.path.join(D, 'panel.pkl'), 'wb'))
for a in ARMS:
    ks = calls[a]; print(a, len(ks), collections.Counter(k[2] for k in ks), collections.Counter(k[1] for k in ks))
