"""2.8 (English edition): build the task set.

  make_tasks.py draft      gemini-3.8-flash reads each skill and proposes 2 candidate tasks in English -> candidates_raw.json
                           (the excerpt must be a verbatim substring of the skill; non-verbatim candidates are dropped)
  make_tasks.py baseline   every candidate, no skill in the window, 3 models, graded by the cross-model judge
                           (lib.PRIMARY_JUDGE) -> candidates_baseline.json, candidates_summary.json
  make_tasks.py finalize   apply the hand review in selection.json -> tasks.json, task_cues.json

Twins (near-duplicate files, where opening either one answers the task) are not used as targets."""
import json, re, sys
from concurrent.futures import ThreadPoolExecutor
from lib import HERE, call, judge, load_pool, PREAMBLE, PRIMARY_JUDGE, CostMeter, append_costs

NOT_TARGET = {'skill_download_paper', 'bestpractice_academic_paper_conversion', 'typefully_post', 'typefully_metrics'}
MODELS = ['deepseek-v4.1-flash', 'gemini-3.8-flash', 'grok-4-fast']
GEN_MODEL = 'gemini-3.8-flash'
meter = CostMeter()

DRAFT_SYS = """You are writing test questions for a controlled experiment. You get one skill document (operating instructions
written for an AI agent). Find 2 rules or facts in it that are specific and non-obvious, for example an exact command-line
flag, a default value, a limit, a step that must not be done, an easy-to-miss pitfall, or a check that must come first.
Then write one question for each.

Requirements:
1. Write the question in English, as a real user's request or question (1 to 3 sentences). It may mention the product or
   domain (for example "my Narwal robot vacuum"), but it must not mention the skill's file name, and it must not copy any
   distinctive sentence or distinctive combination of terms from the document.
2. Only someone who knows this rule can answer correctly; an experienced engineer who has not read this document would
   probably answer wrongly or vaguely. Do not test common knowledge, and do not test private paths on the author's machine.
3. The rule must be stated explicitly in the document, with a single correct answer, not two equally defensible answers.
4. excerpt must be one contiguous passage copied verbatim from the document (do not change a single character,
   30 to 400 characters), and it must support the reference answer on its own.
5. rubric.must_include lists 1 to 3 grading points (specific content the answer must contain); reference_answer is a
   one-sentence reference answer.

Output JSON only: {"tasks":[{"question":"...","excerpt":"...","fact":"one-sentence summary of the rule","rubric":{"must_include":["..."],"reference_answer":"..."}}]}"""


def norm(s):
    return re.sub(r'\s+', ' ', s).strip()


def draft():
    pool, texts = load_pool()
    names = [s['name'] for s in pool['skills'] if s['name'] not in NOT_TARGET]
    out_path = HERE / 'candidates_raw.json'
    done = json.loads(out_path.read_text()) if out_path.exists() else {}

    def one(n):
        if n in done and done[n].get('tasks') is not None:
            return n, done[n]
        r = call(GEN_MODEL, [{'role': 'system', 'content': DRAFT_SYS},
                             {'role': 'user', 'content': f'Skill name: {n}\n\n<document>\n{texts[n]}\n</document>'}],
                 max_tokens=16000)
        meter.add(r)
        if 'error' in r:
            return n, {'error': r['error']}
        txt = r['message'].get('content') or ''
        try:
            d = json.loads(txt[txt.find('{'):txt.rfind('}') + 1])
        except Exception:
            return n, {'error': 'bad json', 'raw': txt}
        kept = []
        for t in d.get('tasks', []):
            t['excerpt_verbatim'] = norm(t.get('excerpt', '')) in norm(texts[n])
            kept.append(t)
        print(n, len(kept), sum(t['excerpt_verbatim'] for t in kept), flush=True)
        return n, {'tasks': kept, 'cost': r['cost']}

    with ThreadPoolExecutor(8) as ex:
        for n, v in ex.map(one, names):
            done[n] = v
            out_path.write_text(json.dumps(done, ensure_ascii=False, indent=1))
    cands = []
    for n in names:
        for i, t in enumerate(done.get(n, {}).get('tasks') or []):
            if t.get('excerpt_verbatim'):
                cands.append({'cid': f'{n}#{i}', 'target': n, **{k: t[k] for k in ('question', 'excerpt', 'fact', 'rubric')}})
    (HERE / 'candidates.json').write_text(json.dumps(cands, ensure_ascii=False, indent=1))
    print('candidates with verbatim excerpt:', len(cands))
    append_costs('2_8_skills/make_tasks.py draft', meter.calls, meter.pt, meter.ct, meter.usd, f'task drafting with {GEN_MODEL}')


def other_skills_with(excerpt, target, texts):
    """Which other skills contain the core of the excerpt (longest line, >=25 chars) verbatim?"""
    lines = sorted((norm(l) for l in excerpt.splitlines()), key=len, reverse=True)
    core = lines[0] if lines else ''
    if len(core) < 25:
        return []
    return [n for n, t in texts.items() if n != target and core in norm(t)]


def baseline():
    _, texts = load_pool()
    cands = json.loads((HERE / 'candidates.json').read_text())
    out_path = HERE / 'candidates_baseline.json'
    done = {(d['cid'], d['model']): d for d in json.loads(out_path.read_text())} if out_path.exists() else {}
    jobs = [(c, m) for c in cands for m in MODELS if (c['cid'], m) not in done or done[(c['cid'], m)].get('correct') is None]

    def one(job):
        c, m = job
        r = call(m, [{'role': 'system', 'content': PREAMBLE}, {'role': 'user', 'content': c['question']}], max_tokens=6000)
        meter.add(r)
        ans = None if 'error' in r else (r['message'].get('content') or '')
        g = judge(c, ans, PRIMARY_JUDGE[m]) if ans is not None else {'correct': None, 'reason': r.get('error'), 'judge_cost': 0}
        meter.usd += g.get('judge_cost') or 0
        return {'cid': c['cid'], 'model': m, 'answer': ans, 'web_tools': r.get('web_tools', []), 'error': r.get('error'),
                'correct': g['correct'], 'reason': g['reason'], 'judge': g.get('judge'), 'cost': r.get('cost'),
                'judge_cost': g.get('judge_cost')}

    res = {k: v for k, v in done.items()}
    with ThreadPoolExecutor(12) as ex:
        for i, d in enumerate(ex.map(one, jobs)):
            res[(d['cid'], d['model'])] = d
            if i % 30 == 0:
                out_path.write_text(json.dumps(list(res.values()), ensure_ascii=False, indent=1))
                print(i, '/', len(jobs), f'${meter.usd:.2f}', flush=True)
    res = list(res.values())
    out_path.write_text(json.dumps(res, ensure_ascii=False, indent=1))
    by = {}
    for d in res:
        by.setdefault(d['cid'], {})[d['model']] = d
    summ = []
    for c in cands:
        b = by.get(c['cid'], {})
        summ.append({'cid': c['cid'], 'target': c['target'],
                     'baseline_correct': sum(1 for m in MODELS if b.get(m, {}).get('correct')),
                     'correct_models': [m for m in MODELS if b.get(m, {}).get('correct')],
                     'used_web': any(b.get(m, {}).get('web_tools') for m in MODELS),
                     'excerpt_also_in': other_skills_with(c['excerpt'], c['target'], texts)})
    (HERE / 'candidates_summary.json').write_text(json.dumps(summ, ensure_ascii=False, indent=1))
    print('all three wrong:', sum(s['baseline_correct'] == 0 for s in summ), '; <=1 right:',
          sum(s['baseline_correct'] <= 1 for s in summ), 'of', len(summ))
    append_costs('2_8_skills/make_tasks.py baseline', meter.calls, meter.pt, meter.ct, meter.usd,
                 'no-skill baseline on all candidates, 3 models + cross-model judge (judge cost included)')


def finalize():
    cands = {c['cid']: c for c in json.loads((HERE / 'candidates.json').read_text())}
    summ = {s['cid']: s for s in json.loads((HERE / 'candidates_summary.json').read_text())}
    sel = json.loads((HERE / 'selection.json').read_text())
    tasks = []
    for i, s in enumerate(sel['selected']):
        c = dict(cands[s['cid']])
        c.update(s.get('edits', {}))
        tasks.append({'id': f'T{i + 1:02d}', 'cid': s['cid'], 'target': c['target'], 'question': c['question'],
                      'excerpt': c['excerpt'], 'fact': c['fact'], 'rubric': c['rubric'], 'cue': s['cue'],
                      'baseline_correct_of_3': summ[s['cid']]['baseline_correct'], 'review_note': s.get('note', '')})
    (HERE / 'tasks.json').write_text(json.dumps(tasks, ensure_ascii=False, indent=1))
    no_cue = [t['id'] for t in tasks if not t['cue']]
    (HERE / 'task_cues.json').write_text(json.dumps(
        {'no_cue': no_cue, 'note': 'Tagged by hand at selection time, before any with-skill run: a question is "no cue" when it '
                                   'names neither the product/tool nor any distinctive term from the target skill\'s name or '
                                   'index line.'}, indent=1))
    print(len(tasks), 'tasks;', len({t['target'] for t in tasks}), 'distinct targets;', len(no_cue), 'no-cue')


if __name__ == '__main__':
    {'draft': draft, 'baseline': baseline, 'finalize': finalize}[sys.argv[1]]()