"""Save per-request outputs and decoding configuration after generation."""
import hashlib
import json
import os
from pathlib import Path
import re


def atomic(path, value):
    path = Path(path)
    temporary = path.with_suffix(path.suffix + '.tmp')
    temporary.write_text(json.dumps(value, ensure_ascii=False, indent=2))
    temporary.replace(path)


def capture(ns):
    root = Path(os.environ['ADSD_RUN_OUTPUT'])
    index = ns['args'].sample_offset + ns['progress'] - 1
    ids = ns['outputs'][0].detach().cpu().tolist()
    prompt = ns['input_ids'][0].detach().cpu().tolist()
    continuation = ids[len(prompt):]
    text = ns['tokenizer1'].decode(continuation, skip_special_tokens=True)
    def last_number(value):
        matches = re.findall(r'-?\d+(?:\.\d+)?', value.replace(',', ''))
        return matches[-1] if matches else None
    counts = ns.get('counts', {}) if ns['args'].speculative else {}
    row = dict(index=index, question=ns['q'], reference=ns['a'],
               prompt_ids=prompt, generated_ids=continuation, generated_text=text,
               elapsed_seconds=ns['end']-ns['start'], counts=counts,
               correct_last_number=last_number(text) == last_number(ns['a']),
               reached_cap=len(continuation) >= ns['args'].max_new_tokens,
               prompt_sha256=hashlib.sha256(json.dumps(prompt).encode()).hexdigest())
    atomic(root / f'row_{index:04d}.json', row)
    atomic(root / 'progress.json', dict(status='running', completed=ns['progress'],
                                       expected=ns['num_samples']))


def policy(ns):
    root = Path(os.environ['ADSD_RUN_OUTPUT'])
    import torch
    import transformers
    import signal
    def interrupted(signum, _frame):
        atomic(root / 'interrupted.json', {'signal': signum,
               'completed_rows': len(list(root.glob('row_*.json'))),
               'recovery': 'Retain partial rows; restart in a new output directory.'})
        raise SystemExit(128 + signum)
    signal.signal(signal.SIGTERM, interrupted)
    signal.signal(signal.SIGINT, interrupted)
    atomic(root / 'effective_policy.json', {
        'torch': torch.__version__, 'transformers': transformers.__version__,
        'models': {name: dict(dtype=str(ns[name].dtype),
                            attention=ns[name].config._attn_implementation,
                            generation_config=ns[name].generation_config.to_dict())
                   for name in ('model1', 'model2')},
        'args': vars(ns['args'])})