"""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}


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

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