"""Condition A: the chat box. One API call with only the user request, then run the returned code once, unmodified.

Usage (from the textbook repo root):
  op run --env-file experiments/aib.env -- experiments/2_1_workspace/.venv/bin/python experiments/2_1_workspace/run_chat.py <model> <trial_no>
"""
import json
import os
import pathlib
import re
import shutil
import subprocess
import sys
import time

HERE = pathlib.Path(__file__).resolve().parent
sys.path.insert(0, str(HERE.parent))
sys.path.insert(0, str(HERE))
from common import chat  # noqa: E402
from check import check  # noqa: E402

REQUEST = ('帮我写一个 Python 脚本，把 memos/ 文件夹里的语音备忘录用 AI Builder Space 的语音识别接口转成文字，'
           '每条存成一个同名的 .txt 文件放在 transcripts/ 里。API key 在环境变量 AI_BUILDER_TOKEN 里。')


def extract_code(text):
    blocks = re.findall(r'```(?:python|py)\s*\n(.*?)```', text, flags=re.S)
    if not blocks:
        blocks = re.findall(r'```[a-zA-Z0-9]*\s*\n(.*?)```', text, flags=re.S)
    return (max(blocks, key=len) if blocks else None), len(blocks)


def main(model, trial):
    tdir = HERE / 'trials' / f'A{trial}'
    if tdir.exists():
        raise SystemExit(f'{tdir} exists; refusing to overwrite')
    tdir.mkdir(parents=True)
    t0 = time.time()
    reply, usage, secs = chat([{'role': 'user', 'content': REQUEST}], model=model)
    (tdir / 'response.md').write_text(reply or '', encoding='utf-8')
    code, n_blocks = extract_code(reply or '')
    run_dir = tdir / 'run'
    shutil.copytree(HERE / 'fixture' / 'memos', run_dir / 'memos')  # fresh copy, no API.md
    result = {'condition': 'A', 'trial': trial, 'model': model, 'chat_seconds': round(secs, 1), 'usage': usage,
              'python_blocks': n_blocks, 'runs': 0}
    if code is None:
        result.update(success=False, failure='no code block in reply')
    else:
        (run_dir / 'script.py').write_text(code, encoding='utf-8')
        env = {k: v for k, v in os.environ.items() if k in ('PATH', 'HOME', 'LANG', 'AI_BUILDER_TOKEN', 'TMPDIR')}
        r0 = time.time()
        try:
            p = subprocess.run([str(HERE / '.venv' / 'bin' / 'python'), 'script.py'], cwd=run_dir, env=env,
                               capture_output=True, text=True, timeout=600)
            rc, out, err = p.returncode, p.stdout, p.stderr
        except subprocess.TimeoutExpired as e:
            rc, out, err = 'timeout', e.stdout or '', e.stderr or ''
        result['runs'] = 1
        result['run_seconds'] = round(time.time() - r0, 1)
        result['exit_code'] = rc
        (tdir / 'run_stdout.txt').write_text(out if isinstance(out, str) else out.decode(), encoding='utf-8')
        (tdir / 'run_stderr.txt').write_text(err if isinstance(err, str) else err.decode(), encoding='utf-8')
        result['check'] = check(run_dir)
        result['success'] = result['check']['success']
        urls = sorted(set(re.findall(r'https?://[^\s\'"`)]+', code)))
        result['urls_in_code'] = urls
    result['total_seconds'] = round(time.time() - t0, 1)
    (tdir / 'result.json').write_text(json.dumps(result, ensure_ascii=False, indent=1), encoding='utf-8')
    print(json.dumps({k: v for k, v in result.items() if k != 'check'}, ensure_ascii=False, indent=1))
    print('check:', result.get('check', {}).get('n_ok'), '/', result.get('check', {}).get('n_memos'))


if __name__ == '__main__':
    main(sys.argv[1], int(sys.argv[2]))