"""2.4: the same cleanup request in a fresh session, with and without the cleanup skill.
Both conditions get the names list (built in 2.3), so the only difference is the skill.
Checks per memo (all must pass for a first-try pass):
  no_um      : no 嗯/啊/呃 left
  keep_tic   : 就是说 kept as many times as the speaker said it
  spacing    : no Chinese character directly touching a Latin letter or digit
  no_rewrite : output length within 85%-115% of the transcript (nothing summarised or merged)"""
import json, pathlib, re, sys
from concurrent.futures import ThreadPoolExecutor
ROOT = pathlib.Path(__file__).resolve().parents[1]; sys.path.insert(0, str(ROOT))
from common import chat, dump
HERE = pathlib.Path(__file__).resolve().parent
d = json.load(open(ROOT / 'corpus/memos.json', encoding='utf-8'))
src = json.load(open(ROOT / 'corpus/transcripts_grok.json', encoding='utf-8'))
ids = [m['id'] for m in d['memos'] if m['set'] in 'WL']
names = (ROOT / '2_3_names/names.md').read_text(encoding='utf-8')
skill = (HERE / 'cleanup_skill.md').read_text(encoding='utf-8')
ref = {m['id']: m['text'] for m in d['memos']}
CJK = r'[一-鿿]'
def check(mid, out):
    spoken_tic = ref[mid].count('就是说')
    return {'no_um': not re.search('[嗯啊呃]', out),
            'keep_tic': out.count('就是说') >= spoken_tic,
            'spacing': not re.search(fr'{CJK}[A-Za-z0-9]|[A-Za-z0-9]{CJK}', out),
            'no_rewrite': 0.85 <= len(out) / max(1, len(src[mid])) <= 1.15}
ASK = '帮我整理一下这段语音转写，只输出整理好的全文。'
conds = {'no_skill': lambda t: f'{ASK}\n\n{names}\n\n转写：\n{t}',
         'with_skill': lambda t: f'{ASK}\n\n{skill}\n\n{names}\n\n转写：\n{t}'}
res = json.load(open(HERE / 'results.json', encoding='utf-8')) if (HERE / 'results.json').exists() else {}
for model in sys.argv[1:] or ['deepseek-v4.1-flash', 'gemini-3.8-flash']:
    for c, mk in conds.items():
        def one(mid):
            out, _, _ = chat([{'role': 'user', 'content': mk(src[mid])}], model, temperature=0)
            return mid, out.strip()
        with ThreadPoolExecutor(8) as ex:
            outs = dict(ex.map(one, ids))
        dump(outs, HERE / f'out_{model}_{c}.json')
        checks = {mid: check(mid, outs[mid]) for mid in ids}
        per_rule = {k: sum(v[k] for v in checks.values()) for k in next(iter(checks.values()))}
        passed = sum(all(v.values()) for v in checks.values())
        res[f'{model}:{c}'] = {'first_try_pass': f'{passed}/{len(ids)}', 'per_rule': per_rule, 'violations': sum(len(ids) - x for x in per_rule.values())}
        print(model, c, res[f'{model}:{c}'])
dump(res, HERE / 'results.json')