"""Execution / agent benchmarks. Execution recipes build a runnable program
(completion + tests + checker) via a harness closure and run it in a sandbox;
agent recipes wait for the agent layer (env_reward slot)."""

from ..recipe import EvalRecipe, register_eval


def _humaneval_harness(sample, pred: str):
    test = sample.metadata.get('test', '')
    entry = sample.metadata.get('entry_point', 'f')
    prog = f'{sample.input}{pred}\n\n{test}\n\ncheck({entry})\nprint("PASSED")\n'
    return {'main.py': prog}


@register_eval('humaneval')
def humaneval():
    return EvalRecipe(
        name='humaneval',
        extract='code_any',
        scorers={'pass': {'name': 'execution', 'harness': _humaneval_harness,
                          'sandbox': 'docker', 'timeout_s': 30}},
        aggregators={'pass': 'pass_at_k'},
        description='HumanEval; completion + official tests in a sandbox, pass@k.',
    )


def _bcb_harness(sample, pred: str):
    test = sample.metadata.get('test', '')
    entry = sample.metadata.get('entry_point', 'f')
    prog = f'{sample.input}{pred}\n\n{test}\n\ncheck({entry})\nprint("PASSED")\n'
    return {'main.py': prog}


@register_eval('bigcodebench')
def bigcodebench():
    return EvalRecipe(
        name='bigcodebench',
        extract='code_any',
        scorers={'pass': {'name': 'execution', 'harness': _bcb_harness,
                          # official image bundles every task's libs (sympy/pandas/...)
                          'image': 'bigcodebench/bigcodebench-eval:latest',
                          'sandbox': 'docker', 'timeout_s': 120}},
        aggregators={'pass': 'pass_at_k'},
        description='BigCodeBench; official all-libs docker image, pass@k.',
    )


_LCB_RUNNER = r'''
import json, subprocess, sys
cases = json.load(open('cases.json'))
failed = 0
for i, case in enumerate(cases):
    stdin = case.get('input', '')
    expected = [str(e).rstrip('\n') for e in ([case['output']] if isinstance(case.get('output'), str) else case.get('output', []))]
    r = subprocess.run([sys.executable, 'solution.py'], input=stdin,
                       capture_output=True, text=True, timeout=20)
    got = [l for l in r.stdout.split('\n') if l != '']
    if got != expected:
        failed += 1
        print(f'case {i}: expected {expected!r} got {got!r}', file=sys.stderr)
if failed:
    print(f'{failed}/{len(cases)} cases failed', file=sys.stderr)
    sys.exit(1)
print('PASSED')
'''


def _lcb_harness(sample, pred: str):
    import json

    starter = sample.metadata.get('starter_code') or ''
    cases = sample.metadata.get('public_test_cases') or '[]'
    cases = json.loads(cases) if isinstance(cases, str) else cases
    return {
        'solution.py': f'{starter}\n{pred}\n',
        'cases.json': json.dumps(cases or []),
        'runner.py': _LCB_RUNNER,
    }


@register_eval('live_code_bench')
def live_code_bench():
    return EvalRecipe(
        name='live_code_bench',
        extract='code_any',
        scorers={'pass': {'name': 'execution', 'harness': _lcb_harness,
                          'entry': 'runner.py', 'sandbox': 'local', 'timeout_s': 60}},
        aggregators={'pass': 'pass_at_k'},
        description='LiveCodeBench; stdin/stdout public-case runner in sandbox.',
    )


import json as _json


def _swe_harness(sample, pred: str):
    """Apply the predicted patch in the official per-instance sweb image and
    run FAIL_TO_PASS (+PASS_TO_PASS) tests. Single-turn protocol: the model
    reads problem_statement and emits a unified diff."""
    f2p = _json.loads(sample.metadata.get('FAIL_TO_PASS') or '[]')
    p2p = _json.loads(sample.metadata.get('PASS_TO_PASS') or '[]')
    tests = f2p + p2p[:20]  # guard: cap regression tests for runtime
    script = f'''set -e
cd /testbed
git apply --whitespace=fix /work/patch.diff || {{ echo PATCH_FAILED; exit 2; }}
FAIL=0
while IFS= read -r t; do
  [ -z "$t" ] && continue
  if ! (conda run -n testbed python -m pytest -x -q "$t" > /dev/null 2>&1); then
    echo "TEST_FAILED $t"; FAIL=1
  fi
done <<'EOF'
{chr(10).join(tests)}
EOF
[ "$FAIL" = 0 ] && echo RESOLVED
exit $FAIL
'''
    return {'patch.diff': pred or '', 'run.sh': script}


@register_eval('swe_bench_verified')
def swe_bench_verified():
    return EvalRecipe(
        name='swe_bench_verified',
        extract='identity',          # a patch, not an answer
        scorers={'resolved': {'name': 'execution', 'harness': _swe_harness,
                              'entry': 'run.sh', 'sandbox': 'docker',
                              'timeout_s': 900}},
        description='SWE-bench Verified single-turn: model emits a unified diff; '
                    'applied in the official sweb.eval.* image, FAIL_TO_PASS(+P2P) '
                    'must pass. Prefetch: evalharness sandbox prefetch swe_bench_verified',
    )


@register_eval('tau2_bench')
def tau2_bench():
    return EvalRecipe(
        name='tau2_bench',
        extract='identity',
        scorers={'acc': 'env_reward'},
        description='tau2-bench; user-simulated dialog, environment reward.',
    )


@register_eval('bfcl_v3')
def bfcl_v3():
    return EvalRecipe(
        name='bfcl_v3',
        extract='identity',
        scorers={'acc': 'env_reward'},  # call-sequence vs ground truth (bfcl_mock env)
        aggregators={'acc': 'weighted_group_avg'},  # group_key = test_category
        description='BFCL v3; run with env=bfcl_mock (agent pump), official call-sequence scoring.',
    )


@register_eval('general_fc')
def general_fc():
    return EvalRecipe(
        name='general_fc',
        extract='identity',
        scorers={'acc': 'execution'},
        description='General function calling; tool-call comparison.',
    )