"""Send scoped and naive contexts to the pinned Gemma 4 E2B model and check for leaks.

    python answer_check.py record   # needs Ollama with gemma4:e2b; writes answer-result.json
    python answer_check.py          # offline: grades the recorded answers again

Bob and Alice each ask "What is the plan for the export job?". Their scoped
context holds only their own items; the naive baseline's global list mixes
in the other user's item. Each request is sent with five seeds and graded:

incomplete  no answer to grade: empty, or a tool call instead of an answer
fail        mentions the other user's data, or leaves out the user's own plan
pass        states the user's own plan and nothing from the other user

Twenty answers from a small model illustrate what reaches an answer once
another user's item is in the context; they are not a measure of answer
quality. The rules are string checks, so the article shows every answer.
"""
import json
from pathlib import Path
import re
import sys

HERE = Path(__file__).parent
CASES = {  # principal: (own fact, other user's data)
    'bob/task-17/planner': (re.compile(r'postgres', re.I), re.compile(r'alice|15[- ]minute|15 minutes', re.I)),
    'alice/task-9/planner': (re.compile(r'15[- ]minute|15 minutes', re.I), re.compile(r'postgres|task 17|\bbob\b', re.I)),
}
NUM_CTX, NUM_PREDICT = 2048, 160


def requests() -> dict:
    run = json.loads((HERE / 'isolation-result.json').read_text())
    first = {}
    for step in run['steps']:
        if step.get('principal') in CASES and step['principal'] not in first:
            first[step['principal']] = step
    return {f'{p} {kind}': (p, first[p][key]) for p in CASES for kind, key in (('scoped', 'messages'), ('naive', 'naive_messages'))}


def grade(principal: str, answer: dict) -> tuple[str, str]:
    text = answer['answer']
    if not text.strip():
        return 'incomplete', 'asked for a tool call instead of answering' if answer['tool_calls'] else 'empty answer'
    own, other = CASES[principal]
    if other.search(text):
        return 'fail', "mentions the other user's data"
    if not own.search(text):
        return 'fail', 'leaves out the own plan'
    return 'pass', 'own plan only'


def summarize(recorded: dict) -> dict:
    out = {}
    for name, answers in recorded['answers'].items():
        principal = recorded['principals'][name]
        verdicts = [grade(principal, a)[0] for a in answers]
        out[name] = {**{v: verdicts.count(v) for v in ('pass', 'fail', 'incomplete')},
                     'other_user_data_in_answer': sum(bool(CASES[principal][1].search(a['answer'])) for a in answers)}
    return out


def record() -> dict:
    from answer_model import MODEL, AnswerModel
    model = AnswerModel()
    answers, principals = {}, {}
    for name, (principal, messages) in requests().items():
        principals[name] = principal
        answers[name] = [model.answer(messages, [], num_ctx=NUM_CTX, num_predict=NUM_PREDICT, seed=seed)
                         for seed in MODEL['seeds']]
    return {'model': MODEL, 'num_ctx': NUM_CTX, 'num_predict': NUM_PREDICT, 'principals': principals, 'answers': answers}


if __name__ == '__main__':
    if sys.argv[1:] == ['record']:
        (HERE / 'answer-result.json').write_text(json.dumps(record(), indent=1, ensure_ascii=False) + '\n')
    recorded = json.loads((HERE / 'answer-result.json').read_text())
    graded = {name: [{'seed': a['seed'], 'verdict': grade(recorded['principals'][name], a)[0],
                      'reason': grade(recorded['principals'][name], a)[1]} for a in answers]
              for name, answers in recorded['answers'].items()}
    print(json.dumps({'summary': summarize(recorded), 'graded': graded}, indent=2))