"""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')
    # BCB official semantics: completion is a standalone module; `test` is a
    # unittest.TestCase subclass -> run it with unittest (official runner uses
    # `unittest.main()` with a buffer; exit 0 == all tests pass)
    prog = f'{pred}\n\n{test}\n\nif __name__ == "__main__":\n    import unittest\n    unittest.main()\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 sandbox image (bundles every task's deps)
                          'image': 'bigcodebench-sandbox: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'))

def as_lines(v):
    """Normalize an expected output to a list of lines (no trailing empties)."""
    if not isinstance(v, list):
        v = [v]
    out = []
    for item in v:
        out.extend(str(item).rstrip('\n').split('\n'))
    return [l for l in out if l != '']

failed = 0
for i, case in enumerate(cases):
    stdin = case.get('input', '')
    expected = as_lines(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_decode_cases(raw):
    """LCB test cases: official data packs private cases as base64+zlib+pickle."""
    import base64
    import io
    import json
    import pickle
    import zlib

    if raw is None:
        return []
    if not isinstance(raw, str):
        return raw if isinstance(raw, list) else []
    try:
        blob = zlib.decompress(base64.b64decode(raw))
        if blob[:2] in (b'\x80\x04', b'\x80\x05', b'\x80\x02'):  # pickle protocol
            data = pickle.load(io.BytesIO(blob))
        else:
            data = json.loads(blob.decode())
    except Exception:
        data = None
    if data is None:
        try:
            data = json.loads(raw)
        except (ValueError, TypeError):
            return []
    # LCB double-packs: pickle list may hold a JSON STRING of the real list
    if isinstance(data, str):
        try:
            data = json.loads(data)
        except (ValueError, TypeError):
            return []
    if isinstance(data, dict):  # {'input':..,'output':..} single case
        data = [data]
    return data if isinstance(data, list) else []


def _lcb_harness(sample, pred: str, use_private: bool = True):
    import json

    starter = sample.metadata.get('starter_code') or ''
    if use_private:
        cases = _lcb_decode_cases(sample.metadata.get('private_test_cases'))
    else:
        cases = _lcb_decode_cases(sample.metadata.get('public_test_cases'))
    if not cases:  # private unavailable -> fall back to public
        cases = _lcb_decode_cases(sample.metadata.get('public_test_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',
    )


def _tau2_reward(pred, target, sample, ctx):
    """Score from the official engine's reward_info (env_state)."""
    env_state = ctx.params.get('env_state') or {}
    rewards = env_state.get('tau2_rewards') or {}
    env_r = rewards.get('environment_reward')
    comm_r = rewards.get('communication_reward')
    vals = [r for r in (env_r, comm_r) if isinstance(r, (int, float))]
    score = float(sum(vals) / len(vals)) if vals else 0.0
    return ({'acc': score}, {'acc': {'mode': 'official_tau2',
                                     'env_reward': env_r,
                                     'comm_reward': comm_r}})


@register_eval('tau2_bench')
def tau2_bench():
    return EvalRecipe(
        name='tau2_bench',
        extract='identity',
        scorers={'acc': _tau2_reward},
        aggregators={'acc': 'grouped_avg'},
        description='tau2-bench via OFFICIAL engine (user simulator + env + reward); '
                    "run with env='tau2_official'",
    )


@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():
    from ..recipe import EvalRecipe, register_eval as _re  # noqa: F401  (keep import local)

    def _gfc_extract(raw, sample):
        # prediction = did the model call any tool? serialized tool_calls in raw
        called = '"name"' in (raw or '') and ('tool_call' in (raw or '').lower()
                                             or raw.strip().startswith('[{"name"'))
        return ('True' if called else 'False'), True, 'tool_called_bool'

    return EvalRecipe(
        name='general_fc',
        extract=_gfc_extract,
        scorers={'acc': {'name': 'exact', 'mode': 'raw'}},
        description='General function calling; predicts should-call-tool (True/False) vs target.',
    )