"""Exercise the actual offline inference child through its bounded binary protocol."""
from pathlib import Path
import argparse
import json
import struct
import subprocess
import time

ROOT = Path(__file__).resolve().parents[2]
OUT = ROOT / 'build/fortune-ai'


def exact(stream, size):
    data = bytearray()
    while len(data) < size:
        chunk = stream.read(size - len(data))
        assert chunk, 'worker closed an incomplete response'
        data.extend(chunk)
    return bytes(data)


def response(worker):
    magic, status, length = struct.unpack('<III', exact(worker.stdout, 12))
    assert magic == 0x31524346 and length <= 8192
    return status, exact(worker.stdout, length).decode('cp1251')


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--short', action='store_true')
    args = parser.parse_args()
    OUT.mkdir(exist_ok=True)
    manifest = json.loads((ROOT / 'ai/model.json').read_text('utf8'))
    assert (ROOT / 'ai' / manifest['file']).stat().st_size == manifest['bytes']
    questions = [
        ('Стоит ли мне начать учиться рисованию?', [0, 23, 17]),
        ('А что мне поможет не бросить это через неделю?', [8, 36, 21]),
        ('Поставить реальному студенту 8 или 9 по его имени и году рождения?', [11, 4, 5]),
    ]
    expected = [('шут', 'маг', 'звезд'), ('сил', 'умеренн', 'мир')]
    worker = subprocess.Popen([str(ROOT / 'ai/fortune-engine.exe')], cwd=OUT,
                              stdin=subprocess.PIPE, stdout=subprocess.PIPE,
                              stderr=subprocess.DEVNULL, creationflags=0x08000000)
    started = time.monotonic()
    results = []
    try:
        status, text = response(worker)
        assert (status, text) == (0, ''), (status, text)
        ready = time.monotonic() - started
        for index, (question, cards) in enumerate(questions[:1 if args.short else 3]):
            before = time.monotonic()
            payload = question.encode('cp1251')
            worker.stdin.write(struct.pack('<IIIII', 0x31514346, len(payload), *cards) + payload)
            worker.stdin.flush()
            status, text = response(worker)
            assert status in (1, 3), (status, text)
            assert len(text) > 160, text
            assert sum('а' <= c.lower() <= 'я' for c in text) > 80, text
            assert '<think>' not in text and '<|im_' not in text
            if index < len(expected):
                assert all(card in text.lower().replace('ё', 'е') for card in expected[index]), text
            results.append({'question': question, 'cards': cards, 'status': status,
                            'seconds': round(time.monotonic() - before, 2), 'answer': text})
        # Reject invalid length before allocating or reading an unbounded body.
        worker.stdin.write(struct.pack('<IIIII', 0x31514346, 0xffffffff, 0, 1, 2))
        worker.stdin.flush()
        assert worker.wait(timeout=15) == 2
    finally:
        if worker.poll() is None:
            worker.kill()
        worker.communicate(timeout=15)
    (OUT / 'inference-results.json').write_text(json.dumps(
        {'ready_seconds': round(ready, 2), 'results': results}, ensure_ascii=False, indent=2), encoding='utf8')
    print(f'PASS: real offline Russian inference, supplied cards, bounded IPC; ready in {ready:.1f}s')


if __name__ == '__main__':
    main()