"""Small real-image smoke test; not a multimodal quality benchmark."""
import argparse
import base64
import io
import json
from pathlib import Path
import httpx
from PIL import Image, ImageDraw


def image_url(color, shape):
    image = Image.new('RGB', (256, 256), 'white')
    draw = ImageDraw.Draw(image)
    getattr(draw, shape)((40, 40, 216, 216), fill=color)
    buffer = io.BytesIO()
    image.save(buffer, format='PNG')
    return 'data:image/png;base64,' + base64.b64encode(buffer.getvalue()).decode()


def validate(base_url):
    red = image_url('red', 'ellipse')
    blue = image_url('blue', 'rectangle')
    def text(value): return {'text': value}
    def picture(value): return {'images': [value]}
    with httpx.Client(base_url=base_url.rstrip('/')+'/', timeout=180) as client:
        health = client.get('health'); health.raise_for_status()
        assert health.json()['scoring'] == 'linear'
        prepared = client.post('materials', json={'material': picture(red)})
        prepared.raise_for_status(); material = prepared.json(); key = material['material_id']
        try:
            questions = [
                {'question': text('What color is the shape in the material?'), 'options': [text('Red'), text('Blue')]},
                {'question': {'text': 'What color is the shape in this question image?', 'images': [blue]}, 'options': [text('Red'), text('Blue')]},
                {'question': text('Which option image matches the material?'), 'options': [picture(blue), picture(red)]},
            ]
            response = client.post(f'materials/{key}/questions', json={'questions': questions})
            response.raise_for_status(); result = response.json()
            assert result['scoring'] == 'linear'
            assert result['question_forward_passes'] == 3
            selected = [a['selected_index'] for a in result['answers']]
            assert selected == [1, 2, 2], selected
            for answer in result['answers']:
                assert answer['question_forward_passes'] == 1
                assert abs(sum(o['probability'] for o in answer['options'])-1) < 1e-6
                assert [o['label'] for o in answer['options']] == ['0', '1']
            return {'health': health.json(), 'material': material, 'result': result, 'expected': [1, 2, 2], 'passed': True}
        finally:
            client.delete(f'materials/{key}').raise_for_status()


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--base-url', default='http://127.0.0.1:8244/jev/api/')
    parser.add_argument('--output', default='multimodal-smoke.json')
    args = parser.parse_args()
    evidence = validate(args.base_url)
    Path(args.output).write_text(json.dumps(evidence, indent=2)+'\n')
    print(json.dumps({'passed': True, 'selected': [a['selected_index'] for a in evidence['result']['answers']], 'scoring': 'linear', 'question_forward_passes': 3}))