"""Local evaluation with resumable answers; never includes a benchmark answer key."""
import argparse
import json
import os
import re
import time
from pathlib import Path
import requests
from huggingface_hub import HfApi
from agent import API, fetch, solve

TOKEN_FILE = Path.home() / '.config' / 'hf-agents-course' / 'token'
OUT = Path(__file__).parent / 'artifacts' / 'results.json'


def format_answer(answer, question):
    answer = str(answer).strip()
    if re.search(r'comma[ -](separated|delimited)', question, re.I):
        answer = ', '.join(part.strip() for part in answer.split(','))
    return answer


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--limit', type=int, default=1)
    parser.add_argument('--submit', action='store_true')
    parser.add_argument('--retry', nargs='*', default=[])
    parser.add_argument('--text-only', action='store_true')
    args = parser.parse_args()
    if not os.getenv('HF_TOKEN') and TOKEN_FILE.is_file():
        os.environ['HF_TOKEN'] = TOKEN_FILE.read_text().strip()
    token = os.getenv('HF_TOKEN')
    if not token:
        raise SystemExit('Configure a Hugging Face inference token first.')
    if HfApi(token=token).whoami()['name'] != 'bestdive':
        raise SystemExit('Token must belong to bestdive.')
    rows = json.loads(OUT.read_text()) if OUT.exists() else []
    if args.submit:
        questions = json.loads(fetch(API + '/questions')[0])
        answers = {r['task_id']: r['submitted_answer'] for r in rows}
        payload = {'username': 'bestdive', 'agent_code': 'https://huggingface.co/spaces/bestdive/agents-course-final/tree/main',
            'answers': [{'task_id': q['task_id'], 'submitted_answer': format_answer(answers.get(q['task_id'], ''),q['question'])} for q in questions]}
        (OUT.parent / 'submission.json').write_text(json.dumps(payload, indent=2))
        r = requests.post(API + '/submit', json=payload, timeout=90)
        r.raise_for_status()
        result = r.json()
        (OUT.parent / 'score.json').write_text(json.dumps(result, indent=2))
        print(json.dumps(result, indent=2))
        return
    done = {r['task_id'] for r in rows}
    questions = json.loads(fetch(API + '/questions')[0])
    attempted = 0
    for question in questions:
        if args.retry:
            if not any(question['task_id'].startswith(prefix) for prefix in args.retry):
                continue
        elif question['task_id'] in done:
            continue
        if args.text_only and (question.get('file_name') or 'youtube.com' in question['question']):
            continue
        if attempted >= max(0, min(args.limit, 20)):
            break
        print('Solving', question['task_id'], flush=True)
        try:
            for attempt in range(3):
                try:
                    answer = solve(question)
                    break
                except Exception as exc:
                    if type(exc).__name__ not in {'ConnectError','ConnectTimeout','ReadTimeout','ConnectionError'} or attempt == 2:
                        raise
                    print('Retrying transient network connection',attempt + 1,flush=True)
                    time.sleep(2 * (attempt + 1))
        except Exception as exc:
            raise SystemExit('Stopped after model/tool failure: ' + type(exc).__name__ + '; checkpoint preserved.')
        rows = [r for r in rows if r['task_id'] != question['task_id']]
        rows.append({'task_id': question['task_id'], 'question': question['question'], 'submitted_answer': answer})
        OUT.parent.mkdir(parents=True, exist_ok=True)
        temporary = OUT.with_suffix('.tmp')
        temporary.write_text(json.dumps(rows, ensure_ascii=False, indent=2))
        temporary.replace(OUT)
        attempted += 1
        print('Saved result', len(rows), flush=True)


if __name__ == '__main__':
    main()