#!/usr/bin/env python3
"""Scorer for BENCHMARK_COUNTERPARTY_V1, rules fixed before any held-out run:
- correct: decision equals truth AND failed_checks equals the truth set exactly
- critical_error: a false approval (APPROVE_PAYMENT while truth is HOLD) or a truth-failed check the agent did not list
- false_approval: APPROVE_PAYMENT while truth is HOLD, counted separately as the costliest error
- unjustified_hold: HOLD while truth is APPROVE_PAYMENT, a failure, never safety
- no answer after at most 2 failed attempts: incorrect, with the full time of all attempts
- cost per run in EUR: 0.01 per wall second of all attempts + 2 EUR per million reported tokens (CLI total, lower bound); unknown tokens reported, never guessed
- primary: cost per correct completion and critical errors per task, per condition; 95 percent bootstrap intervals over tasks
- pre-declared target: C_ASTRANL at least 30 percent fewer critical errors per task than C_WEB, or at least 2x lower cost per correct
  completion, with a completion rate not lower than C_WEB
Usage: python3 score.py [--results DIR] [--paired-only]"""
import argparse, json, random
from pathlib import Path
HERE = Path('/opt/astranl/var/benchmark_counterparty_v1')


def judge(r, tr):
    tf = set(tr['failed_checks'])
    if not r.get('ok'):
        return {'correct': False, 'critical': 1 + len(tf), 'false_approval': 0, 'unjustified_hold': 0, 'no_answer': True, 'missed': len(tf)}
    a = r['answer']; d = a.get('decision'); f = set(a.get('failed_checks') or [])
    fa = int(d == 'APPROVE_PAYMENT' and tr['decision'] == 'HOLD')
    missed = len(tf - f)
    return {'correct': d == tr['decision'] and f == tf, 'critical': fa + missed, 'false_approval': fa,
            'unjustified_hold': int(d == 'HOLD' and tr['decision'] == 'APPROVE_PAYMENT'), 'no_answer': False, 'missed': missed}


def cost(r):
    sec = (r.get('seconds') or 0) + sum((x.get('seconds') or 0) for x in r.get('failed_attempts', []) if r.get('ok'))
    tok = r.get('tokens_total')
    return 0.01 * sec + (tok * 2.0 / 1e6 if tok else 0.0), tok is None


def summ(rows):
    n = len(rows); c = sum(x['j']['correct'] for x in rows); tot = sum(x['cost'] for x in rows)
    return {'runs': n, 'correct': c, 'completion_rate': round(c / n, 4) if n else None, 'cost_eur_total': round(tot, 3),
            'cost_per_correct_eur': round(tot / c, 4) if c else None,
            'critical_errors_per_task': round(sum(x['j']['critical'] for x in rows) / n, 4) if n else None,
            'false_approvals': sum(x['j']['false_approval'] for x in rows), 'unjustified_holds': sum(x['j']['unjustified_hold'] for x in rows),
            'no_answer': sum(x['j']['no_answer'] for x in rows), 'tokens_unknown_runs': sum(x['unk'] for x in rows),
            'astranl_tool_used_runs': sum(1 for x in rows if x['mcp']), 'median_seconds': sorted(x['sec'] for x in rows)[n // 2] if n else None}


def boot(by, cond, key, B=2000):
    ids = sorted(by); R = random.Random(428); vals = []
    for _ in range(B):
        s = [R.choice(ids) for _ in ids]
        rr = [by[i][cond] for i in s if cond in by[i]]
        if key == 'critical':
            vals.append(sum(x['j']['critical'] for x in rr) / max(1, len(rr)))
        else:
            c = sum(x['j']['correct'] for x in rr); vals.append(sum(x['cost'] for x in rr) / c if c else float('inf'))
    vals.sort()
    return [round(vals[int(0.025 * B)], 4), round(vals[int(0.975 * B) - 1], 4)]


def main():
    ap = argparse.ArgumentParser(); ap.add_argument('--results', default=str(HERE / 'results')); ap.add_argument('--truth', default=str(HERE / 'truth.json'))
    ap.add_argument('--paired-only', action='store_true')
    a = ap.parse_args()
    truth = {t['task_id']: t for t in json.load(open(a.truth))['truth']}
    by = {}
    for f in Path(a.results).glob('*.json'):
        r = json.load(open(f))
        if r['task_id'] not in truth:
            continue
        if not r.get('ok') and len(r.get('failed_attempts', [])) < 2:
            continue
        c, unk = cost(r)
        by.setdefault(r['task_id'], {})[r['condition']] = {'j': judge(r, truth[r['task_id']]), 'cost': c, 'unk': unk, 'sec': r.get('seconds') or 0,
                                                          'mcp': bool((r.get('trace') or {}).get('mcp_calls')), 'cat': truth[r['task_id']]['category']}
    if a.paired_only:
        by = {k: v for k, v in by.items() if len(v) == 2}
    out = {'tasks_scored': len(by), 'conditions': {}, 'by_category_correct': {}}
    for cond in ('C_WEB', 'C_ASTRANL'):
        rows = [v[cond] for v in by.values() if cond in v]
        out['conditions'][cond] = summ(rows)
        if rows:
            out['conditions'][cond]['critical_ci95'] = boot(by, cond, 'critical')
            out['conditions'][cond]['cost_per_correct_ci95'] = boot(by, cond, 'cost')
        for cat in sorted({x['cat'] for x in rows}):
            out['by_category_correct'].setdefault(cat, {})[cond] = '%d/%d' % (sum(x['j']['correct'] for x in rows if x['cat'] == cat), sum(1 for x in rows if x['cat'] == cat))
    print(json.dumps(out, indent=1))


if __name__ == '__main__':
    main()