"""Verify real server CPU fallback when CUDA is requested but unavailable.

Required arguments: --server (compiled binary), --library (CPU-only native
ORT), --model (export directory), --model-id and --out (new receipt directory).
Optional --threads defaults to 4. No parent TOKENLING_* configuration is used.
The child uses a disposable private database and a random loopback port.
Saves/access keys stay in memory or the disposable DB, never in its receipt.
"""
import argparse
import json
import os
import secrets
import socket
import subprocess
import time
import urllib.error
import urllib.request
import uuid
from pathlib import Path


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    for key in ('server', 'library', 'model', 'model-id', 'out'):
        parser.add_argument('--' + key, required=True)
    parser.add_argument('--threads', type=int, default=4)
    args = parser.parse_args()
    if args.threads < 1:
        parser.error('threads must be positive')
    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=False)
    with socket.socket() as selected:
        selected.bind(('127.0.0.1', 0))
        port = selected.getsockname()[1]
    base = f'http://127.0.0.1:{port}'
    env = {k: v for k, v in os.environ.items() if not k.startswith('TOKENLING_')}
    env.update(TOKENLING_LISTEN=f'127.0.0.1:{port}', TOKENLING_DB=str(out / 'private.sqlite'),
               TOKENLING_MODEL_DIR=args.model, TOKENLING_MODEL_ID=args.model_id,
               TOKENLING_ALLOWED_ORIGIN=base, TOKENLING_THREADS=str(args.threads),
               TOKENLING_GPU_DEVICE='0', ORT_DYLIB_PATH=args.library, RAYON_NUM_THREADS='1')

    def api(path, key=None, body=None, method=None):
        headers = {'Origin': base}
        if key:
            headers['Authorization'] = 'Bearer ' + key
        if body is not None:
            headers['Content-Type'] = 'application/json'
        request = urllib.request.Request(base + '/api/' + path,
            data=json.dumps(body).encode() if body is not None else None,
            headers=headers, method=method)
        with urllib.request.urlopen(request, timeout=10) as response:
            assert response.headers['Cache-Control'] == 'no-store'
            assert response.headers['Access-Control-Allow-Origin'] == base
            return json.load(response)

    with (out / 'server.stdout.log').open('w') as stdout, (out / 'server.stderr.log').open('w') as stderr:
        server = subprocess.Popen([args.server], env=env, stdout=stdout, stderr=stderr)
        try:
            deadline = time.monotonic() + 30
            while True:
                assert server.poll() is None, 'server exited during startup'
                try:
                    health = api('health')
                    break
                except urllib.error.URLError:
                    assert time.monotonic() < deadline, 'server startup timed out'
                    time.sleep(.05)
            assert health['cpu'] and not health['gpu'], 'CPU-only library did not fall back'
            state = json.loads(Path(__file__).with_name('fixtures').joinpath('save-v1.json').read_text())
            now = time.time() * 1000
            state.update(v=2, hosting={'mode': 'local', 'brainId': args.model_id, 'timeZone': 'UTC'},
                         created=now, lastTick=now, lastActive=now, nextPoop=now + 86400000,
                         wokeAt=now, grimeSince=0, asleep=False, sick=False, poop=0, offer=None,
                         chat=['Sam: hello little friend'], history={})
            state['pet'].update(hatched=now, stage='baby')
            state['settings']['chatter'] = False
            state['stats'] = {'tummy': 60, 'fun': 60, 'energy': 80, 'clean': 80}
            pet_id, key, device = str(uuid.uuid4()), secrets.token_hex(32), str(uuid.uuid4())
            api('pets', key, {'id': pet_id, 'state': state, 'modelId': args.model_id})
            try:
                api('pets/' + pet_id)
                raise AssertionError('private state was readable without a key')
            except urllib.error.HTTPError as error:
                assert error.code == 401
            api('pets/' + pet_id + '/transfer', key,
                {'epoch': 0, 'operationId': str(uuid.uuid4()), 'mode': 'remote', 'deviceId': device, 'state': state})
            deadline = time.monotonic() + 30
            while True:
                snapshot = api('pets/' + pet_id, key)
                assert snapshot['lastError'] is None, 'CPU fallback generation failed'
                if not snapshot['thinking'] and snapshot['state']['chat'] != state['chat']:
                    break
                assert time.monotonic() < deadline, 'CPU fallback generation timed out'
                time.sleep(.02)
            parked = api('pets/' + pet_id + '/transfer', key,
                {'epoch': 1, 'operationId': str(uuid.uuid4()), 'mode': 'local', 'deviceId': device, 'state': None})
            assert parked['authority'] == 'local'
            api('pets/' + pet_id, key, method='DELETE')
            (out / 'receipt.json').write_text(json.dumps({'modelId': args.model_id,
                'cpuThreads': args.threads, 'gpuRequested': True, 'gpuAvailable': False,
                'realCpuReply': True, 'privateAccess': True, 'corsAndNoStore': True,
                'localTransferAndDelete': True}, indent=2) + '\n')
        finally:
            server.terminate()
            try:
                server.wait(timeout=10)
            except subprocess.TimeoutExpired:
                server.kill()
                server.wait()


if __name__ == '__main__':
    main()