"""Pure source-role and prompt rules for the prospective prefix-teacher pilot."""
import hashlib
import json
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
ARMS = ('native', 'empty_context', 'solution_context')
TEMPLATE = ('Here is an earlier solution to the problem. Use it to continue the solution below.\n'
            '<earlier_solution>\n{solution}\n</earlier_solution>\n\n{native_prompt}')


def sha(path):
    h = hashlib.sha256()
    with Path(path).open('rb') as stream:
        for block in iter(lambda: stream.read(8 * 1024 * 1024), b''):
            h.update(block)
    return h.hexdigest()


def save_new(path, value):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open('x') as stream:
        stream.write(json.dumps(value, indent=2) + '\n')


def select_roles(rows, eos, prefix_tokens=64):
    """Never consult the prefix draw's outcome. Missing roles remain explicit."""
    by_sample = {row['sample']: row for row in rows}
    assert len(by_sample) == len(rows)
    assert all(sample in by_sample for sample in range(5))
    teacher = next((by_sample[s] for s in range(4)
                    if by_sample[s]['correct'] and by_sample[s]['stop'] == 'eos'), None)
    prefix = by_sample[4]['token_ids'][:prefix_tokens]
    reasons = []
    if teacher is None:
        reasons.append('no_verified_eos_context_in_draws_0_to_3')
    if len(prefix) < prefix_tokens or eos in prefix:
        reasons.append('prefix_draw_ended_before_unended_64_tokens')
    return teacher, prefix, reasons


def prompts(tokenizer, native_prompt, prefix, teacher_ids):
    assert teacher_ids[-1] == tokenizer.eos_token_id
    assert tokenizer.eos_token_id not in teacher_ids[:-1]
    solution = tokenizer.decode(teacher_ids[:-1])
    texts = {'native': native_prompt,
             'empty_context': TEMPLATE.format(solution='', native_prompt=native_prompt),
             'solution_context': TEMPLATE.format(solution=solution, native_prompt=native_prompt)}
    # Append the exact original prefix IDs. Never decode and retokenize the prefix.
    return {arm: tokenizer.encode(text) + prefix for arm, text in texts.items()}
