"""Drive engine/host.mjs: the game runs in Node, the scene comes back as arrays."""
import json
import os
import struct
import subprocess
from pathlib import Path

import numpy as np

ROOT = Path(__file__).resolve().parent
EDEN = Path(os.environ.get('EDEN_DEPLOY_ROOT', '/root/autodl-tmp/edenlive'))
GAMES = Path(os.environ.get('EDEN_GAMES_ROOT', EDEN / 'EdenLive/realtime/games'))
NODE = EDEN / 'node/bin/node'

_DTYPES = {'Float32Array': np.float32, 'Uint32Array': np.uint32, 'Uint16Array': np.uint16, 'Int32Array': np.int32}


def standalone_config(entry):
    """The same window.gameStandalone payload the capture browser's page gets."""
    return {'restartKeys': list(entry.get('restart_keys') or ['r']),
            'eventPrompts': dict(entry.get('event_prompts') or {})}


class Engine:
    def __init__(self, slug='eiffel-assault', width=960, height=544, games_root=None, config=None):
        games = Path(games_root or GAMES)
        if config is None:
            entry = next(g for g in json.loads((games / 'registry.json').read_text())['games'] if g['slug'] == slug)
            config = standalone_config(entry)
        self.proc = subprocess.Popen(
            [str(NODE), str(ROOT / 'engine/host.mjs'), str(games / slug),
             str(games / 'node_modules/three/build/three.module.js'), str(width), str(height), json.dumps(config)],
            stdin=subprocess.PIPE, stdout=subprocess.PIPE, bufsize=0)
        self.scene, self.scene_arrays = self._read()
        if self.scene.get('type') != 'scene':
            raise RuntimeError(f'Unexpected first message: {self.scene.get("type")}')

    def _read_exact(self, n):
        chunks, remaining = [], n
        while remaining:
            chunk = self.proc.stdout.read(remaining)
            if not chunk:
                raise EOFError('Engine exited')
            chunks.append(chunk)
            remaining -= len(chunk)
        return b''.join(chunks)

    def _read(self):
        (size,) = struct.unpack('<I', self._read_exact(4))
        header = json.loads(self._read_exact(size))
        (size,) = struct.unpack('<I', self._read_exact(4))
        binary = self._read_exact(size)
        arrays = {name: np.frombuffer(binary, _DTYPES[spec['dtype']], spec['length'], spec['offset'])
                  for name, spec in header.pop('arrays', {}).items()}
        return header, arrays

    def advance(self, target_tick, events=(), capture=False, state=True):
        self.send_advance(target_tick, events, capture, state)
        return self.receive()

    def send_advance(self, target_tick, events=(), capture=False, state=True):
        """Queue one step without waiting; replies come back in order through receive()."""
        command = {'cmd': 'advance', 'targetTick': int(target_tick), 'events': list(events),
                   'capture': bool(capture), 'state': bool(state)}
        self.proc.stdin.write((json.dumps(command) + '\n').encode())

    def receive(self):
        return self._read()

    def send_chunk(self, frames):
        """frames: [{'targetTick', 'events', 'capture'}]. Read back with receive():
        one ('scene', arrays) message per captured frame, then the closing 'chunk'."""
        command = {'cmd': 'chunk', 'frames': [dict(targetTick=int(f['targetTick']), events=list(f['events']),
                                                   capture=bool(f['capture'])) for f in frames]}
        self.proc.stdin.write((json.dumps(command) + '\n').encode())

    def probe(self, events=()):
        self.proc.stdin.write((json.dumps({'cmd': 'probe', 'events': list(events)}) + '\n').encode())
        return self._read()[0]

    def close(self):
        if self.proc.poll() is None:
            try:
                self.proc.stdin.write(b'{"cmd":"quit"}\n')
                self.proc.wait(5)
            except Exception:
                self.proc.kill()