"""2.8 (English edition): run one model under one or more conditions over tasks.json. Raw outputs go to raw/<model>__<condition>.json.

  run.py <model> <cond> [<cond> ...] [--only T01,T05]
  conditions: none, full_1, full_10, full_30, full_all, index, index_rep2, index_fixed

Identical design to the Chinese run (experiments/2_7_skills/run.py), with English prompts:
full_k : the system prompt holds k full skills: the target plus the first k-1 non-target skills of one fixed random
         permutation of the pool (seed 27); the target goes to a seeded random slot. full_all is the whole pool in permutation order.
index  : the system prompt holds the index of all skills (name + one-line description, 240-char cap); the model opens skills,
         at most 4 opens, then must answer. grok uses native tool calls (read_skill). deepseek and gemini cannot use client
         tools through the gateway, so they use a text protocol: a line `READ_SKILL: <name>`, answered by a user message
         carrying the skill text.
index_rep2 : an exact repeat of index, to measure run-to-run variation at temperature 0 (also the control for index_fixed).
index_fixed : index with descriptions_fixed.json overriding some one-line descriptions; runs only the task ids listed for
         that model in index_fixed_plan.json."""
import json, random, re, sys, time
from concurrent.futures import ThreadPoolExecutor
from lib import HERE, PREAMBLE, call, load_pool, index_text, skills_block, CostMeter, append_costs

SEED = 27
MAX_OPENS = 4
RAW = HERE / 'raw'
RAW.mkdir(exist_ok=True)

pool, texts = load_pool()
NAMES = [s['name'] for s in pool['skills']]
PERM = NAMES[:]
random.Random(SEED).shuffle(PERM)
TASKS = json.loads((HERE / 'tasks.json').read_text())
meter = CostMeter()

TOOLS = [{'type': 'function', 'function': {
    'name': 'read_skill', 'description': 'Open one skill and return its full text.',
    'parameters': {'type': 'object', 'properties': {'name': {'type': 'string', 'description': 'skill name exactly as in the index'}},
                   'required': ['name']}}}]


def full_set(task, k):
    if k >= len(NAMES):
        return PERM[:], PERM.index(task['target'])
    others = [n for n in PERM if n != task['target']][:k - 1]
    pos = random.Random(f'{task["id"]}-{k}').randint(0, k - 1)
    return others[:pos] + [task['target']] + others[pos:], pos


def max_tok(model):
    return 16000 if model.startswith('gemini') else 4000


def one_call(model, msgs, **kw):
    r = call(model, msgs, max_tokens=max_tok(model), **kw)
    meter.add(r)
    return r


def run_full(model, task, k):
    names, pos = full_set(task, k)
    sys_p = f'{PREAMBLE}\n\nBelow are the skill documents you can refer to, {len(names)} in total:\n\n{skills_block(names, texts)}'
    r = one_call(model, [{'role': 'system', 'content': sys_p}, {'role': 'user', 'content': task['question']}])
    u = r.get('usage') or {}
    return {'answer': None if 'error' in r else (r['message'].get('content') or ''), 'error': r.get('error'),
            'n_skills': len(names), 'target_pos': pos, 'prompt_tokens': u.get('prompt_tokens'),
            'completion_tokens': u.get('completion_tokens'), 'sec': r.get('sec'), 'cost': r.get('cost'),
            'web_tools': r.get('web_tools', []), 'n_calls': 1, 'finish_reason': r.get('finish_reason'), 'nudged': r.get('nudged')}


def run_none(model, task):
    r = one_call(model, [{'role': 'system', 'content': PREAMBLE}, {'role': 'user', 'content': task['question']}])
    u = r.get('usage') or {}
    return {'answer': None if 'error' in r else (r['message'].get('content') or ''), 'error': r.get('error'),
            'prompt_tokens': u.get('prompt_tokens'), 'completion_tokens': u.get('completion_tokens'), 'sec': r.get('sec'),
            'cost': r.get('cost'), 'web_tools': r.get('web_tools', []), 'n_calls': 1, 'nudged': r.get('nudged')}


def index_system(overrides, text_protocol):
    how = ('When you need the full text of a skill, call the tool read_skill(name) to open it. You may open several, but at '
           'most 4 in total. After opening, answer the user.'
           if not text_protocol else
           'When you need the full text of a skill, reply with only this line: READ_SKILL: <name> (the name must match the '
           'index exactly; you may write several such lines, one per line; at most 4 in total), and write nothing else in '
           'that reply. The system will send you the full text, and then you answer the user. If you do not need to open '
           'anything, answer directly.')
    return (f'{PREAMBLE}\n\nYou have a skill library. Below is its index; each line is a skill name and a one-line '
            f'description:\n\n{index_text(pool, overrides)}\n\n{how}')


def read_skill(name):
    name = (name or '').strip().strip('`"\' ')
    if name in texts:
        return name, texts[name]
    return None, f'There is no skill named {name!r}. Use a name from the index.'


def run_index(model, task, overrides):
    text_protocol = not model.startswith('grok')
    msgs = [{'role': 'system', 'content': index_system(overrides, text_protocol)}, {'role': 'user', 'content': task['question']}]
    opened, requested, calls = [], [], []
    answer, err = None, None
    for step in range(MAX_OPENS + 2):
        capped = len(requested) >= MAX_OPENS
        if text_protocol:
            r = one_call(model, msgs)
        else:
            r = one_call(model, msgs, tools=TOOLS, tool_choice='none' if capped else None)
        calls.append({k: r.get(k) for k in ('usage', 'sec', 'cost', 'web_tools', 'error', 'finish_reason', 'nudged')})
        if 'error' in r:
            err = r['error']
            break
        m = r['message']
        if text_protocol:
            content = m.get('content') or ''
            reqs = re.findall(r'^\s*\**READ_SKILL:\**\s*`?([^\s`*]+)`?\s*$', content, re.M)
            if reqs and not capped:
                msgs.append({'role': 'assistant', 'content': content})
                parts = []
                for q in reqs:
                    if len(requested) >= MAX_OPENS:
                        parts.append(f'[read_skill {q}] The limit of 4 opens has been reached; not opened.')
                        continue
                    requested.append(q)
                    n, body = read_skill(q)
                    if n:
                        opened.append(n)
                    parts.append(f'[read_skill {q}]\n{body}')
                tail = '\n\nThe open limit has been reached. Answer the user directly now.' if len(requested) >= MAX_OPENS else ''
                msgs.append({'role': 'user', 'content': '\n\n'.join(parts) + tail})
                continue
            answer = content
            break
        tcs = m.get('tool_calls') or []
        if tcs and not capped:
            msgs.append({k: v for k, v in m.items() if v is not None})
            for tc in tcs:
                try:
                    q = json.loads(tc['function'].get('arguments') or '{}').get('name')
                except Exception:
                    q = None
                if len(requested) >= MAX_OPENS:
                    body = 'The limit of 4 opens has been reached; not opened. Answer the user directly now.'
                else:
                    requested.append(q)
                    n, body = read_skill(q)
                    if n:
                        opened.append(n)
                msgs.append({'role': 'tool', 'tool_call_id': tc['id'], 'content': body})
            continue
        answer = m.get('content') or ''
        break
    pt = sum((c['usage'] or {}).get('prompt_tokens') or 0 for c in calls)
    ct = sum((c['usage'] or {}).get('completion_tokens') or 0 for c in calls)
    return {'answer': answer, 'error': err, 'opened': opened, 'requested': requested,
            'opened_target': task['target'] in opened, 'prompt_tokens': pt, 'completion_tokens': ct,
            'sec': sum(c['sec'] or 0 for c in calls), 'cost': sum(c['cost'] or 0 for c in calls),
            'web_tools': [w for c in calls for w in (c['web_tools'] or [])], 'n_calls': len(calls),
            'nudged': any(c.get('nudged') for c in calls)}


def main():
    args = sys.argv[1:]
    only = None
    if '--only' in args:
        i = args.index('--only')
        only = set(args[i + 1].split(','))
        args = args[:i] + args[i + 2:]
    model, conds = args[0], args[1:]
    for cond in conds:
        tasks = TASKS
        overrides = {}
        if cond == 'index_fixed':
            overrides = json.loads((HERE / 'descriptions_fixed.json').read_text())['fixed']
            plan = json.loads((HERE / 'index_fixed_plan.json').read_text())[model]
            ids = set(plan['failed_not_opened'] + plan['passing_sample'])
            tasks = [t for t in TASKS if t['id'] in ids]
        if only:
            tasks = [t for t in tasks if t['id'] in only]
        out_path = RAW / f'{model}__{cond}.json'
        prev = {d['id']: d for d in json.loads(out_path.read_text())} if out_path.exists() else {}
        todo = [t for t in tasks if t['id'] not in prev or prev[t['id']].get('error') or prev[t['id']].get('answer') is None]

        def one(t):
            if cond == 'none':
                d = run_none(model, t)
            elif cond.startswith('full_'):
                k = cond.split('_')[1]
                d = run_full(model, t, len(NAMES) if k == 'all' else int(k))
            else:
                d = run_index(model, t, overrides)
            return {'id': t['id'], 'target': t['target'], 'model': model, 'condition': cond, **d}

        workers = 3 if cond == 'full_all' else (6 if cond == 'full_30' else 10)
        t0 = time.time()
        usd0 = meter.usd
        with ThreadPoolExecutor(workers) as ex:
            for d in ex.map(one, todo):
                prev[d['id']] = d
                out_path.write_text(json.dumps(sorted(prev.values(), key=lambda x: x['id']), ensure_ascii=False, indent=1))
                print(model, cond, d['id'], 'err' if d.get('error') else 'ok', d.get('prompt_tokens'), round(d.get('sec') or 0, 1),
                      ('opened=' + ','.join(d['opened'])) if 'opened' in d else '', (d.get('error') or '')[:200], flush=True)
        print('DONE', model, cond, len(todo), 'tasks', round(time.time() - t0), 's', f'${meter.usd - usd0:.3f}', flush=True)
    append_costs(f'2_8_skills/run.py {model} {" ".join(conds)}', meter.calls, meter.pt, meter.ct, meter.usd,
                 'gateway estimated_cost_usd per call (trace)')


if __name__ == '__main__':
    main()