"""2.7: prepare the index_fixed condition.

  fix_descriptions.py plan      from graded index runs: per model, task ids where the answer was wrong AND the target was
                                not opened, plus a seeded sample of 10 previously-correct tasks -> index_fixed_plan.json
  fix_descriptions.py rewrite   for every target skill in any model's failed list, gpt-5 rewrites the one-line description
                                from the skill text alone (it never sees the tasks) -> descriptions_fixed.json
The rewritten descriptions were then read by hand before running index_fixed (see README)."""
import json, random, sys
from concurrent.futures import ThreadPoolExecutor
from lib import HERE, call, load_pool
from grade import load_joined

REWRITE_SYS = """下面是一份给 AI agent 用的 skill 文档和它在索引里的一句话说明。索引里只放这一句话，agent 靠它决定要不要打开全文。
请重写这句话（中文为主，可以保留英文专有名词，不超过 200 个字符），写清楚：什么情况下应该打开它（用户会提出什么样的需求、碰到什么问题），
它覆盖的对象（产品、工具、命令）。不要列出文档里的具体参数值或答案，只描述何时使用。只输出这句话本身。"""


def plan():
    rows = [r for r in load_joined() if r['condition'] == 'index']
    out = {}
    for model in sorted({r['model'] for r in rows}):
        rs = sorted([r for r in rows if r['model'] == model], key=lambda r: r['id'])
        failed = [r['id'] for r in rs if not r['correct'] and not r.get('opened_target')]
        passing = [r['id'] for r in rs if r['correct']]
        sample = sorted(random.Random(f'fix-{model}').sample(passing, min(10, len(passing))))
        out[model] = {'failed_not_opened': failed, 'passing_sample': sample,
                      'failed_targets': sorted({r['target'] for r in rs if r['id'] in failed}),
                      'failed_with_target_opened': [r['id'] for r in rs if not r['correct'] and r.get('opened_target')]}
    (HERE / 'index_fixed_plan.json').write_text(json.dumps(out, ensure_ascii=False, indent=1))
    print(json.dumps(out, ensure_ascii=False, indent=1))


def rewrite():
    pool, texts = load_pool()
    old = {s['name']: s['description'] for s in pool['skills']}
    p = json.loads((HERE / 'index_fixed_plan.json').read_text())
    targets = sorted({t for v in p.values() for t in v['failed_targets']})

    def one(n):
        r = call('gpt-5', [{'role': 'system', 'content': REWRITE_SYS},
                           {'role': 'user', 'content': f'原说明：{old[n]}\n\n<document name="{n}">\n{texts[n]}\n</document>'}],
                 max_tokens=8000)
        return n, (r['message'].get('content') or '').strip() if 'error' not in r else None

    with ThreadPoolExecutor(8) as ex:
        fixed = dict(ex.map(one, targets))
    out = {'fixed': {k: v for k, v in fixed.items() if v}, 'before': {k: old[k] for k in fixed}}
    (HERE / 'descriptions_fixed.json').write_text(json.dumps(out, ensure_ascii=False, indent=1))
    for k in fixed:
        print(k, '\n  BEFORE:', old[k][:240], '\n  AFTER :', fixed[k])


if __name__ == '__main__':
    {'plan': plan, 'rewrite': rewrite}[sys.argv[1]]()