"""2.8 (English edition): prepare the index_fixed condition. Same procedure as experiments/2_7_skills/fix_descriptions.py.

  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, gemini-3.8-flash rewrites the one-line
                                description from the skill text alone (it never sees the tasks) -> descriptions_fixed.json
The rewritten descriptions are 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, CostMeter, append_costs
from grade import load_joined

WRITER = 'gemini-3.8-flash'
REWRITE_SYS = """Below is a skill document written for an AI agent, and its one-line description from the index. The index
holds only this one line, and the agent decides from it alone whether to open the full text.
Rewrite this line (in English, at most 200 characters). State clearly when it should be opened (what the user will ask for,
or what problem they run into) and what it covers (products, tools, commands). Do not list specific parameter values or
answers from the document; describe only when to use it. Output only the line itself."""


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']})
    meter = CostMeter()

    def one(n):
        r = call(WRITER, [{'role': 'system', 'content': REWRITE_SYS},
                          {'role': 'user', 'content': f'Current line: {old[n]}\n\n<document name="{n}">\n{texts[n]}\n</document>'}],
                 max_tokens=8000)
        meter.add(r)
        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 = {'writer': WRITER, 'fixed': {k: v for k, v in fixed.items() if v}, 'before': {k: old[k] for k in fixed},
           'hand_edits': {}}
    (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])
    append_costs('2_8_skills/fix_descriptions.py rewrite', meter.calls, meter.pt, meter.ct, meter.usd, f'writer {WRITER}')


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