#!/usr/bin/env python3
"""Send one prompt set to a chat server, and grade the saved answers.

    compare.py run LABEL URL MODEL OUT.json [ID-PREFIX]
    compare.py grade OUT.json ...
    compare.py selftest
"""
import json
import os
import re
import subprocess
import sys
import tempfile
import time
import urllib.request

import prompts


def chat(url, model, prompt, max_tokens):
    """Ask once, without thinking, at temperature 0."""
    body = {'model': model, 'temperature': 0, 'seed': 1,
            'max_tokens': max_tokens, 'reasoning_effort': 'none',
            'think': False, 'thinking': {'type': 'disabled'},
            'chat_template_kwargs': {'enable_thinking': False},
            'thinking_budget_tokens': 0,
            'messages': [{'role': 'user', 'content': prompt}]}
    request = urllib.request.Request(
        url, json.dumps(body).encode(),
        {'Content-Type': 'application/json'})
    start = time.time()
    with urllib.request.urlopen(request, timeout=7200) as response:
        reply = json.load(response)
    message = reply['choices'][0]['message']
    usage = reply.get('usage', {})
    return dict(text=message.get('content') or '',
                reasoning=len(message.get('reasoning_content') or ''),
                finish=reply['choices'][0].get('finish_reason'),
                prompt_tokens=usage.get('prompt_tokens'),
                completion_tokens=usage.get('completion_tokens'),
                seconds=round(time.time() - start, 1))


def run(label, url, model, out, prefix=''):
    results = {}
    if os.path.exists(out):
        results = json.load(open(out))['results']
    for p in prompts.prompts():
        if not p['id'].startswith(prefix) or p['id'] in results:
            continue
        try:
            r = chat(url, model, p['prompt'], p['max_tokens'])
        except Exception as e:
            r = dict(error=repr(e)[:300])
        results[p['id']] = r
        print(label, p['id'], r.get('seconds'), r.get('prompt_tokens'),
              r.get('completion_tokens'), r.get('error', ''), flush=True)
        with open(out + '.tmp', 'w') as f:
            json.dump(dict(label=label, model=model, results=results), f,
                      indent=1)
        os.replace(out + '.tmp', out)


def code_block(text):
    blocks = re.findall(r'```(?:python|py)?[ \t]*\n(.*?)```', text, re.S)
    return blocks[-1] if blocks else text


def passes_test(code, test):
    """Run the answer's code with the test in a separate process."""
    with tempfile.TemporaryDirectory() as tmp:
        path = os.path.join(tmp, 'answer.py')
        with open(path, 'w') as f:
            f.write(code + '\n# TESTS\n' + test)
        try:
            done = subprocess.run([sys.executable, '-I', path], cwd=tmp,
                                  capture_output=True, timeout=20)
        except subprocess.TimeoutExpired:
            return False
        return done.returncode == 0


def final_answer(text):
    parts = re.split(r'ANSWER\s*:', text)
    if len(parts) < 2:
        return None
    lines = parts[-1].strip().splitlines()
    return lines[0].strip(' *`"\'.') if lines else ''


def answer_matches(text, want):
    got = final_answer(text)
    if got is None:
        return False
    if isinstance(want, str):
        return got.casefold() == want.casefold()
    # "EIP-155" answers a question about an EIP number.
    number = re.match(r'-?[\d,]+(\.\d+)?', re.sub(r'^EIP-', '', got))
    if not number:
        return False
    try:
        return float(number.group().replace(',', '')) == want
    except ValueError:
        return False


def json_matches(text, want):
    try:
        return json.loads(text.strip()) == want
    except ValueError:
        return False


def grade_one(p, r):
    if 'error' in r:
        return False
    if p['kind'] == 'code':
        return passes_test(code_block(r['text']), p['test'])
    if p['kind'] == 'answer':
        return answer_matches(r['text'], p['want'])
    return json_matches(r['text'], p['want'])


def grade(paths):
    table = {}
    for path in paths:
        data = json.load(open(path))
        marks, rate = {}, []
        for p in prompts.prompts():
            r = data['results'].get(p['id'])
            if r is None:
                continue
            if p['kind'] == 'open':
                if r.get('completion_tokens') and r.get('seconds'):
                    rate.append(r['completion_tokens'] / r['seconds'])
                continue
            marks[p['id']] = grade_one(p, r)
        table[data['label']] = dict(
            marks=marks, think=sum(1 for r in data['results'].values()
                                   if r.get('reasoning')),
            rate=round(sum(rate) / len(rate), 1) if rate else None)
    labels = list(table)
    groups = ['code', 'math', 'json', 'find-30k', 'find-100k',
              'eth-fact', 'eth-rub']
    print('%-12s' % 'group', *('%-22s' % x for x in labels))
    for g in groups:
        cells = []
        for x in labels:
            m = [v for k, v in table[x]['marks'].items() if k.startswith(g)]
            cells.append('%-22s' % ('%d/%d' % (sum(m), len(m))))
        print('%-12s' % g, *cells)
    print('%-12s' % 'total', *('%-22s' % ('%d/%d' % (
        sum(table[x]['marks'].values()), len(table[x]['marks'])))
        for x in labels))
    print('%-12s' % 'tokens/s', *('%-22s' % table[x]['rate']
                                  for x in labels))
    print('%-12s' % 'thought', *('%-22s' % table[x]['think']
                                 for x in labels))
    for x in labels:
        print(x, 'failed:', ' '.join(k for k, v in table[x]['marks'].items()
                                     if not v) or 'none')
    json.dump(table, open('grades.json', 'w'), indent=1)


def selftest():
    """Check that each reference solution passes and a stub fails."""
    for p in prompts.prompts():
        if p['kind'] == 'code':
            assert passes_test(p['ref'], p['test']), p['id']
            assert not passes_test('', p['test']), p['id']
        elif p['kind'] == 'answer':
            assert answer_matches('So.\n**ANSWER: %s**' % p['want'],
                                  p['want']), p['id']
            assert not answer_matches('ANSWER: 1x', p['want']), p['id']
        elif p['kind'] == 'json':
            assert json_matches(json.dumps(p['want']), p['want']), p['id']
    print('selftest passed,', len(prompts.prompts()), 'prompts')


if __name__ == '__main__':
    if sys.argv[1] == 'run':
        run(*sys.argv[2:])
    elif sys.argv[1] == 'grade':
        grade(sys.argv[2:])
    else:
        selftest()