"""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')
    base = (sample.metadata or {}).get('prompt') or sample.input
    prog = f'{base}{pred}\n\n{test}\n\ncheck({entry})\nprint("PASSED")\n'
    return {'main.py': prog}


def _humaneval_extract(raw, sample):
    # es/official contract asks for 'ONLY the code' -> the model emits a bare
    # function with no markdown fence; fall back to the raw text then
    from ..extractor import make_extractor
    val, ok, note = make_extractor('code_any')(raw, sample)
    if ok:
        return val, ok, note
    body = (raw or '').strip()
    if body:
        return body, True, 'bare_code'
    return '', False, 'empty'


@register_eval('humaneval')
def humaneval():
    return EvalRecipe(
        name='humaneval',
        extract=_humaneval_extract,
        scorers={'pass': {'name': 'execution', 'harness': _humaneval_harness,
                          'sandbox': 'docker', 'timeout_s': 30}},
        aggregators={'pass': 'pass_at_k'},
        exec_workers=8,
        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);
                          # its ENTRYPOINT is the official evaluate CLI which
                          # swallows our runner -> override with plain python3
                          'image': 'bigcodebench/bigcodebench-evaluate:latest',
                          'entrypoint': 'python3',
                          'sandbox': 'docker', 'timeout_s': 120}},
        aggregators={'pass': 'pass_at_k'},
        exec_workers=12,
        description='BigCodeBench; official all-libs docker image, pass@k.',
    )


_LCB_RUNNER = r'''
import json, subprocess, sys
cases = json.load(open('cases.json'))
meta = json.load(open('meta.json')) if __import__('os').path.exists('meta.json') else {}
fn_name = meta.get('fn_name')

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
if fn_name:
    # function-call style (LeetCode / starter_code problems, es-official):
    # import the solution and call fn_name on each input, compare to output
    import importlib.util
    spec = importlib.util.spec_from_file_location('solution', 'solution.py')
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    fn = getattr(mod, fn_name, None)
    if fn is None:
        # starter classes: instantiate and look for the method on the class
        for attr in vars(mod).values():
            if isinstance(attr, type) and hasattr(attr, fn_name):
                fn = getattr(attr(), fn_name)
                break
    if fn is None:
        print(f'fn_name {fn_name!r} not found in solution', file=sys.stderr)
        sys.exit(1)
    for i, case in enumerate(cases):
        try:
            raw_in, raw_out = case['input'], case['output']
            # lite packs fn-style args/results as JSON STRINGS
            args = json.loads(raw_in) if isinstance(raw_in, str) else raw_in
            expected = json.loads(raw_out) if isinstance(raw_out, str) else raw_out
            args = args if isinstance(args, list) else [args]
            got = fn(*args)
        except Exception as e:
            print(f'case {i}: raised {type(e).__name__}: {e}', file=sys.stderr)
            failed += 1
            continue
        expected = tuple(expected) if isinstance(expected, list) else expected
        got_t = tuple(got) if isinstance(got, list) else got
        if got_t != expected:
            print(f'case {i}: expected {expected!r} got {got_t!r}', file=sys.stderr)
            failed += 1
else:
    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 ''
    # es-official case composition: PUBLIC + PRIVATE in full (use_private
    # toggles the private half; es load_utils.py always uses both)
    pub = _lcb_decode_cases(sample.metadata.get('public_test_cases'))
    priv = _lcb_decode_cases(sample.metadata.get('private_test_cases')) if use_private else []
    cases = pub + priv
    if not cases:
        cases = _lcb_decode_cases(sample.metadata.get('public_test_cases')) or []
    files = {
        'solution.py': f'{starter}\n{pred}\n',
        'cases.json': json.dumps(cases or []),
        'runner.py': _LCB_RUNNER,
    }
    fn_name = (sample.metadata.get('fn_name')
               or _lcb_fn_name_from_metadata(sample.metadata.get('raw_metadata')))
    if fn_name:
        files['meta.json'] = json.dumps({'fn_name': fn_name})
    return files


def _lcb_fn_name_from_metadata(raw):
    """Official lite packs fn_name inside the record's `metadata` JSON blob."""
    import json

    if not raw:
        return None
    try:
        md = json.loads(raw) if isinstance(raw, str) else raw
        return md.get('func_name')
    except Exception:
        return None


@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'},
        exec_workers=8,
        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}


def _official_swebench_eval(instance_meta: dict, patch: str, timeout: int = 1800):
    """Score ONE instance via the OFFICIAL swebench harness -- the exact
    code path es uses (make_test_spec -> build_container -> apply patch ->
    eval script -> get_eval_report). Requires pip install swebench."""
    from pathlib import Path
    import tempfile

    from swebench.harness.grading import get_eval_report
    from swebench.harness.test_spec.test_spec import make_test_spec

    test_spec = make_test_spec(instance_meta, namespace='swebench',
                               arch='x86_64')
    log_dir = Path(tempfile.mkdtemp(prefix='eh-swe-eval-')) / test_spec.instance_id
    log_dir.mkdir(parents=True)

    from docker.client import DockerClient

    client = DockerClient.from_env()
    # official image must exist locally (pull handled by our preflight)
    image_name = test_spec.instance_image_key
    container = client.containers.run(
        image_name, command='tail -f /dev/null', detach=True,
        name=f'eh-swe-official-{test_spec.instance_id[:40]}',
        # the official eval_script runs `pip install -e .[test]`; pypi.org
        # is unreachable from this host's containers -> same tuna mirror
        # as the agent containers
        environment={'PIP_INDEX_URL':
                     'https://pypi.tuna.tsinghua.edu.cn/simple'})
    try:
        patch_file = log_dir / 'patch.diff'
        patch_file.write_text(patch or '')
        container.exec_run(f'cp {patch_file} /tmp/patch.diff' if False
                           else 'mkdir -p /tmp', workdir='/testbed')
        # copy patch in
        import tarfile
        import io as _io

        with tempfile.NamedTemporaryFile(suffix='.tar') as tf:
            with tarfile.open(tf.name, 'w') as tar:
                info = tarfile.TarInfo('patch.diff')
                data = (patch or '').encode()
                info.size = len(data)
                tar.addfile(info, _io.BytesIO(data))
            with open(tf.name, 'rb') as f:
                container.put_archive('/tmp', f.read())
        applied = False
        for cmd in ('git apply -v /tmp/patch.diff',
                    'git apply --reject /tmp/patch.diff',
                    'patch -p1 -i /tmp/patch.diff'):
            r = container.exec_run(cmd, workdir='/testbed')
            if r.exit_code == 0:
                applied = True
                break
        if not applied:
            return {'resolved': 0.0, 'error': 'patch_apply_failed'}
        eval_file = log_dir / 'eval.sh'
        eval_file.write_text(test_spec.eval_script)
        with tempfile.NamedTemporaryFile(suffix='.tar') as tf:
            with tarfile.open(tf.name, 'w') as tar:
                info = tarfile.TarInfo('eval.sh')
                data = test_spec.eval_script.encode()
                info.size = len(data)
                tar.addfile(info, _io.BytesIO(data))
            with open(tf.name, 'rb') as f:
                container.put_archive('/tmp', f.read())
        container.exec_run('chmod +x /tmp/eval.sh')
        r = container.exec_run('/bin/bash /tmp/eval.sh', workdir='/testbed')
        # exec_run without demux gives combined output in .output
        test_output = r.output.decode('utf-8', 'replace')
        (log_dir / 'test_output.txt').write_text(test_output)
        report = get_eval_report(
            test_spec=test_spec,
            prediction={'model_patch': patch, 'instance_id': test_spec.instance_id},
            test_log_path=log_dir / 'test_output.txt',
            include_tests_status=True,
        )
        return {'resolved': float(report.get('resolved', 0.0)),
                'report': {k: v for k, v in report.items()
                           if isinstance(v, (int, float, str, bool))}}
    finally:
        try:
            container.remove(force=True)
        except Exception:
            pass


def _swe_official_reward(pred, target, sample, ctx):
    """Agentic + single-turn shared scorer: official harness when the
    swebench package is importable, our in-container protocol otherwise."""
    es = ctx.params.get('env_state') or {}
    patch = es.get('patch') or pred or ''
    if not patch or '(no patch produced)' in patch:
        return ({'acc': 0.0}, {'acc': {'error': 'no patch'}})
    # metadata from ENV_STATE, not sample.metadata: the runner shuffles
    # samples and restores predictions by checkpoint key, so positional
    # pairing can misalign -- env_state carries the ACTUAL instance the
    # agent ran (verified: sample.metadata said 12907 while env_state
    # held 13453's test_patch, building an eval_script for the wrong bug)
    md = dict(es)
    md.update({k: v for k, v in (sample.metadata or {}).items()
               if k not in md or not md.get(k)})
    # official make_test_spec needs the raw instance fields
    for k in ('FAIL_TO_PASS', 'PASS_TO_PASS'):
        v = md.get(k)
        if isinstance(v, str):
            try:
                md[k] = _json.loads(v)
            except Exception:
                pass
    md.setdefault('repo', str(md.get('repo', '')))
    md.setdefault('version', str(md.get('version', '')))
    try:
        from swebench.harness.test_spec.test_spec import make_test_spec  # noqa

        have = True
    except ImportError:
        have = False
    if have:
        try:
            res = _official_swebench_eval(md, patch)
            return ({'acc': res.get('resolved', 0.0)},
                    {'acc': {'mode': 'official_swebench',
                             'resolved': res.get('resolved'),
                             'error': res.get('error', ''),
                             'turns_used': es.get('turns_used')}})
        except Exception as e:
            return ({'acc': 0.0}, {'acc': {'mode': 'official_swebench',
                                           'error': f'{type(e).__name__}: {str(e)[:150]}'}})
    # fallback: our in-container protocol (no swebench package)
    return _swe_agentic_reward(pred, target, sample, ctx)


@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 _swe_agentic_harness(sample, pred: str):
    """Agentic scoring: pred is the recovered patch (from env_state); same
    in-container test protocol as the single-turn variant."""
    md = sample.metadata or {}
    es = md.get('env_state') or {}
    patch = es.get('patch') or pred or ''
    f2p = _json.loads(es.get('FAIL_TO_PASS') or '[]') if isinstance(
        es.get('FAIL_TO_PASS'), str) else (es.get('FAIL_TO_PASS') or [])
    p2p = _json.loads(es.get('PASS_TO_PASS') or '[]') if isinstance(
        es.get('PASS_TO_PASS'), str) else (es.get('PASS_TO_PASS') or [])
    test_patch = es.get('test_patch') or md.get('test_patch') or ''
    tests = f2p + p2p[:20]
    script = f'''set -e
cd /testbed
cat > /work/model_patch.diff << 'EOFP'
{patch}
EOFP
cat > /work/test_patch.diff << 'EOFT'
{test_patch}
EOFT
git apply --whitespace=fix /work/model_patch.diff || {{ echo PATCH_FAILED; exit 2; }}
git apply --whitespace=fix /work/test_patch.diff || {{ echo TESTPATCH_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 {'run.sh': script}


@register_eval('swe_bench_verified_agentic')
def swe_bench_verified_agentic():
    return EvalRecipe(
        name='swe_bench_verified_agentic',
        extract='identity',
        scorers={'resolved': {'name': 'env_reward',
                              'backend': _swe_official_reward}},
        exec_workers=2,
        description='SWE-bench Verified AGENTIC (mini-swe-agent protocol): '
                    'multi-turn bash agent explores /testbed in the '
                    'per-instance container, submits a git diff; scored by '
                    'FAIL_TO_PASS(+P2P) with the official test patch applied.',
    )


def _tau2_reward(pred, target, sample, ctx):
    """Score from the official engine's reward_info (env_state).

    Current tau2 reward_info carries the COMPOSITE 'reward' plus detail
    fields (db_check / action_checks / communicate_checks); the old
    environment_reward/communication_reward split no longer exists."""
    env_state = ctx.params.get('env_state') or {}
    rewards = env_state.get('tau2_rewards') or {}
    r = rewards.get('reward')
    if not isinstance(r, (int, float)):
        vals = [v for v in (rewards.get('environment_reward'),
                            rewards.get('communication_reward'))
                if isinstance(v, (int, float))]
        r = sum(vals) / len(vals) if vals else 0.0
    return ({'acc': float(r)}, {'acc': {'mode': 'official_tau2',
                                        'reward': r,
                                        'db_check': rewards.get('db_check'),
                                        'note': str((rewards.get('info') or {}).get('note', ''))[:120]}})


@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.',
    )