"""Rewrites the Chinese gloss of every bank word and gives it an example
sentence, using the AI service configured in assets/config/ai_config.json.

The downloaded word books are a dictionary dump: some entries carry only a rare
sense (`ticket` as a verb and nothing else), some are unusable as a quiz answer.
Cleaning rules cannot invent the missing sense, so the model rewrites the gloss
outright, with the book's version passed only as a hint it may overrule.

Results are cached one JSON object per line, so an interrupted run resumes
instead of paying for the same words twice.
"""

import json
import os
import subprocess
import sys
import tempfile
from concurrent.futures import ThreadPoolExecutor

ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
BANK = os.path.join(ROOT, 'assets/words/wordbank.json')
CONFIG = os.path.join(ROOT, 'assets/config/ai_config.json')

# A word the model reasons about at length can truncate its whole batch; set
# ENRICH_BATCH smaller to pick those stragglers up on a second run.
BATCH = int(os.environ.get('ENRICH_BATCH', '12'))
WORKERS = 8
RETRIES = 3

SYSTEM = """你在为一个中国成年人用的英语学习 App 整理词库。对每个词输出：

- zh：最常用的那一个中文释义。2 到 6 个字，不要词性标注，不要括号，不要罗列同义词。
- more：其他常用义，最多 2 个，用「；」分隔；没有别的常用义就给空字符串。
- ex：一个英文例句。日常口语场景，6 到 12 个词，必须用到 zh 所指的那个义项，
  句中其他词不要超过这个词的 CEFR 等级。
- ex_zh：例句的中文翻译，说人话，不要翻译腔。

输入里的「参考」来自一份质量不稳定的词书：有的只收了罕见义，有的缺了最常用的义项。
不要盲从，以你自己对这个词最常用义的判断为准。

只输出 JSON 数组，每个元素形如
{"id":"...","zh":"...","more":"...","ex":"...","ex_zh":"..."}
不要输出任何别的文字，不要用 markdown 代码块。"""


def load_config():
    with open(CONFIG) as handle:
        return json.load(handle)


def call(config, prompt):
    """One chat completion, through curl: this machine's Python does not trust
    the local certificate chain but curl does."""
    body = {
        'model': config['model'],
        'messages': [
            {'role': 'system', 'content': SYSTEM},
            {'role': 'user', 'content': prompt},
        ],
        'temperature': 0.3,
        # The configured model reasons before it answers, and those tokens come
        # out of the same budget -- a tight cap truncates the JSON mid-array.
        'max_tokens': 8000,
    }
    with tempfile.NamedTemporaryFile('w', suffix='.json', delete=False) as payload:
        json.dump(body, payload, ensure_ascii=False)
        path = payload.name
    # The key goes through stdin, never argv, so it stays out of the process list.
    conf = (
        'url = "%s/v1/chat/completions"\n'
        'header = "Authorization: Bearer %s"\n'
        'header = "Content-Type: application/json"\n'
        'data-binary = "@%s"\n'
        'max-time = 300\n'
    ) % (config['endpoint'].rstrip('/'), config['apiKey'], path)
    try:
        done = subprocess.run(
            ['curl', '-sS', '--config', '-'], input=conf, capture_output=True, text=True
        )
    finally:
        os.unlink(path)
    if done.returncode != 0:
        raise RuntimeError(done.stderr[:200])
    answer = json.loads(done.stdout)
    if 'choices' not in answer:
        raise RuntimeError(done.stdout[:200])
    choice = answer['choices'][0]
    if choice.get('finish_reason') == 'length':
        raise RuntimeError('truncated at max_tokens')
    return choice['message']['content'], answer.get('usage', {})


def parse(text):
    """The model is told to return bare JSON; strip a code fence when it adds one."""
    text = text.strip()
    if text.startswith('```'):
        text = text.split('\n', 1)[1].rsplit('```', 1)[0]
    start, end = text.find('['), text.rfind(']')
    if start < 0 or end < 0:
        raise ValueError(text[:200])
    return json.loads(text[start : end + 1])


def ask(config, batch):
    lines = [
        '%s | %s | %s | 参考：%s'
        % (word['id'], word['en'], word['level'], '；'.join(
            filter(None, [word.get('zh', ''), word.get('more', '')])) or '（无）')
        for word in batch
    ]
    text, usage = call(config, '\n'.join(lines))
    wanted = {word['id'] for word in batch}
    rows = [row for row in parse(text) if row.get('id') in wanted]
    return rows, usage


def main():
    limit = int(sys.argv[1]) if len(sys.argv) > 1 else 0
    cache_path = os.environ.get('ENRICH_CACHE', os.path.join(tempfile.gettempdir(), 'enrich_cache.jsonl'))
    done = {}
    if os.path.exists(cache_path):
        with open(cache_path) as handle:
            for line in handle:
                row = json.loads(line)
                done[row['id']] = row
    with open(BANK) as handle:
        bank = json.load(handle)
    todo = [word for word in bank['words'] if word['id'] not in done]
    if limit:
        todo = todo[:limit]
    print('cached %d, to do %d' % (len(done), len(todo)), flush=True)

    config = load_config()
    batches = [todo[i : i + BATCH] for i in range(0, len(todo), BATCH)]
    spent = {'in': 0, 'out': 0, 'fail': 0}
    with open(cache_path, 'a') as cache:
        with ThreadPoolExecutor(max_workers=WORKERS) as pool:
            for index, result in enumerate(pool.map(lambda b: safe(config, b), batches)):
                rows, usage, error = result
                if error:
                    spent['fail'] += 1
                    print('  batch %d failed: %s' % (index, error), flush=True)
                    continue
                for row in rows:
                    cache.write(json.dumps(row, ensure_ascii=False) + '\n')
                cache.flush()
                spent['in'] += usage.get('prompt_tokens', 0)
                spent['out'] += usage.get('completion_tokens', 0)
                if index % 10 == 0:
                    print('  %d/%d batches' % (index + 1, len(batches)), flush=True)
    print('tokens in %d out %d, failed batches %d' % (spent['in'], spent['out'], spent['fail']))
    print('cache: %s' % cache_path)


def safe(config, batch):
    """Retries a batch a few times; a truncated or malformed answer is common
    enough that giving up on the first one would leave holes everywhere."""
    last = 'no attempt'
    for attempt in range(RETRIES):
        try:
            rows, usage = ask(config, batch)
            if rows:
                return rows, usage, None
            last = 'empty result'
        except Exception as error:
            last = '%s: %s' % (type(error).__name__, error)
    return [], {}, last


if __name__ == '__main__':
    main()