"""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
import threading

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):
    deadline = threading.Timer(180, worker.kill)
    deadline.start()
    try:
        magic, status, length = struct.unpack('<III', exact(worker.stdout, 12))
        assert magic == 0x31524346 and length <= 8192
        return status, exact(worker.stdout, length).decode('cp1251')
    finally:
        deadline.cancel()


def check_reading(text, cards):
    catalogue = (ROOT / 'ai/cards.txt').read_text('utf8').splitlines()
    names = [entry.split(' — ')[0] for entry in catalogue]
    roles = ['Ситуация', 'Препятствие', 'Направление']
    lines = text.split('\n')
    assert len(lines) == 4, text
    for line, card, role in zip(lines, cards, roles):
        heading = f'{role} - {names[card % 22]} (' + ('перевёрнутая' if card >= 22 else 'прямая') + '): '
        assert line.startswith(heading), (heading, line)
        body = line[len(heading):]
        assert 80 <= len(body) <= 220, body
        assert body[-1] in '.!?', body
    assert lines[3].startswith('Общий совет: '), text
    assert '<' not in text and '*' not in text, text
    return lines


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]),
    ]
    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 == 1, (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
            check_reading(text, cards)
            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()