#!/usr/bin/env python3
"""Paced, resumable runner for BENCHMARK_COUNTERPARTY_V1. Same agent (ChatGPT subscription via codex exec, live web search,
read-only sandbox), same prompt, same schema in both conditions:
  C_WEB      web search only: the sources a team without AstraNL has
  C_ASTRANL  web search plus the public AstraNL MCP server; the prompt does not name AstraNL, the tool list does
No API key, no paid API. Never reads truth.json. Stops at once on a subscription limit.
Usage: python3 runner.py --max-runs N [--sleep S]"""
import argparse, json, os, random, re, subprocess, sys, tempfile, time
from pathlib import Path
sys.path.insert(0, '/opt/astranl/organism')
import openai_verifier as ov

HERE = Path('/opt/astranl/var/benchmark_counterparty_v1')
CHECKS = ['registered_active', 'name_matches_register', 'vat_valid', 'vat_belongs_to_company']
SCHEMA = {'type': 'object', 'additionalProperties': False, 'required': ['decision', 'failed_checks', 'sources_used', 'note'],
          'properties': {'decision': {'type': 'string', 'enum': ['APPROVE_PAYMENT', 'HOLD']},
                         'failed_checks': {'type': 'array', 'items': {'type': 'string', 'enum': CHECKS}, 'maxItems': 4},
                         'sources_used': {'type': 'array', 'items': {'type': 'string', 'maxLength': 200}, 'maxItems': 12},
                         'note': {'type': 'string', 'maxLength': 600}}}
PROMPT = ('The file ./invoice_case.json holds one supplier invoice, a payment policy and a question. Read it and answer the question '
          'with the JSON object of the schema. Use only information you actually obtained from sources during this task and from the file; '
          'when you could not verify a check, treat it as failed. failed_checks uses only the ids in check_ids.')
MCP = ['-c', 'mcp_servers.astranl.url="https://astranl.com/mcp/streamable"']


def order(tasks, seed=428):
    pairs = [(t['task_id'], c) for t in tasks for c in ('C_WEB', 'C_ASTRANL')]
    random.Random(seed).shuffle(pairs)
    return pairs


def tokens(text):
    m = re.findall(r'tokens used[:\s]+([\d,]+)', text, re.I)
    return int(m[-1].replace(',', '')) if m else None


def run_one(task, cond, timeout=420):
    model_id = ov.models()[0]
    with tempfile.TemporaryDirectory(prefix='astranl-cpb-') as work:
        os.chmod(work, 0o700)
        Path(work, 'invoice_case.json').write_text(json.dumps(task, ensure_ascii=False, indent=1), encoding='utf-8')
        schema = Path(work, 'schema.json'); out = Path(work, 'answer.json')
        schema.write_text(json.dumps(SCHEMA), encoding='utf-8')
        cmd = [ov.CODEX_BIN, '--search'] + (MCP if cond == 'C_ASTRANL' else []) + [
            'exec', '--model', model_id, '--sandbox', 'read-only', '--ephemeral', '--ignore-user-config', '--skip-git-repo-check',
            '--output-schema', str(schema), '--output-last-message', str(out), '--color', 'never', '--cd', work, '-']
        t0 = time.time()
        try:
            p = subprocess.run(cmd, input=PROMPT, capture_output=True, text=True, timeout=timeout, cwd=work, env=ov._subscription_env())
        except subprocess.TimeoutExpired:
            return {'ok': False, 'reason': 'timeout', 'seconds': round(time.time() - t0, 2), 'model_id': model_id}
        sec = round(time.time() - t0, 2)
        log = p.stdout + p.stderr
        low = log.lower()
        ans = ov._parse_json(out.read_text(encoding='utf-8')) if out.exists() else None
        trace = {'web_searches': low.count('web search:'), 'mcp_calls': re.findall(r'mcp: (astranl/[a-z_]+) (?:started)', low),
                 'mcp_failed': low.count('(failed)')}
        if p.returncode != 0 or not ans:
            reason = 'subscription_limit' if any(x in low for x in ('usage limit', 'rate limit', 'limit reached')) else 'not_logged_in' if 'not logged in' in low else 'cli_error_or_schema'
            return {'ok': False, 'reason': reason, 'seconds': sec, 'model_id': model_id, 'exit': p.returncode, 'trace': trace}
        return {'ok': True, 'answer': ans, 'seconds': sec, 'model_id': model_id, 'tokens_total': tokens(log), 'trace': trace}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--max-runs', type=int, default=1); ap.add_argument('--sleep', type=int, default=360)
    ap.add_argument('--tasks', default=str(HERE / 'tasks.json')); ap.add_argument('--out', default=str(HERE / 'results'))
    a = ap.parse_args()
    if 'truth' in os.path.basename(a.tasks):
        sys.exit('refused: the runner never reads a truth file')
    tasks = {t['task_id']: t for t in json.load(open(a.tasks))['tasks']}
    outd = Path(a.out); outd.mkdir(exist_ok=True)
    done = 0
    for tid, cond in order(list(tasks.values())):
        f = outd / ('%s_%s.json' % (tid, cond))
        if f.exists() and json.load(open(f)).get('ok'):
            continue
        if done >= a.max_runs:
            break
        prev = json.load(open(f)) if f.exists() else {}
        if len(prev.get('failed_attempts', [])) >= 2:
            continue
        r = run_one(tasks[tid], cond)
        r.update({'task_id': tid, 'condition': cond, 'at': time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime()),
                  'failed_attempts': prev.get('failed_attempts', []) + ([] if r.get('ok') else [{'reason': r.get('reason'), 'seconds': r.get('seconds')}])})
        f.write_text(json.dumps(r, ensure_ascii=False, indent=1), encoding='utf-8')
        print(json.dumps({k: r.get(k) for k in ('task_id', 'condition', 'ok', 'reason', 'seconds', 'tokens_total', 'trace')}), flush=True)
        done += 1
        if not r.get('ok') and r.get('reason') in ('subscription_limit', 'not_logged_in'):
            print('STOP ' + r['reason'], flush=True)
            break
        if done < a.max_runs:
            time.sleep(a.sleep)


if __name__ == '__main__':
    main()